Repository: tenstorrent/tt-metal
Pull request #20529: [tt-train] Forward pass for CrossEntropyLoss
Author: login: vmelnykovTT
Created: 2025-04-10T18:40:12Z | Merged: 2025-04-30T20:28:12Z

### Ticket
[[Feature Request] Cross entropy loss forward and backward functions](https://github.com/tenstorrent/tt-metal/issues/14113)

### Problem description
Currently, a cross entropy loss operation does not exist.
Instead, we use a composite operation to perform both the forward and backward computations manually.

### What's changed
In this PR, the `cross_entropy_fw` operation is added.
The `cross_entropy_bw` operation will be implemented in a separate PR

### Checklist
- [ ] [All post commit](https://github.com/tenstorrent/tt-metal/actions/workflows/all-post-commit-workflows.yaml) CI passes
- [ ] [Blackhole Post commit](https://github.com/tenstorrent/tt-metal/actions/workflows/blackhole-post-commit.yaml) CI passes (if applicable)
- [ ] [Model regression](https://github.com/tenstorrent/tt-metal/actions/workflows/perf-models.yaml) CI passes (if applicable)
- [ ] [Device performance regression](https://github.com/tenstorrent/tt-metal/actions/workflows/perf-device-models.yaml) CI passes (if applicable)
- [ ] **(For models and ops writers)** Full [new models tests](https://github.com/tenstorrent/tt-metal/actions/workflows/full-new-models-suite.yaml) CI passes (if applicable)
- [x] New/Existing tests provide coverage for changes


Diff:
diff --git a/tt-train/sources/ttml/CMakeLists.txt b/tt-train/sources/ttml/CMakeLists.txt
index 89cda269d8be..eee71de14fb2 100644
--- a/tt-train/sources/ttml/CMakeLists.txt
+++ b/tt-train/sources/ttml/CMakeLists.txt
@@ -39,6 +39,14 @@ set(METAL_OPS_FILES
     ${CMAKE_CURRENT_SOURCE_DIR}/metal/ops/rmsnorm_fw/device/rmsnorm_fw_device_operation.cpp
     ${CMAKE_CURRENT_SOURCE_DIR}/metal/ops/rmsnorm_fw/device/rmsnorm_fw_program_factory.hpp
     ${CMAKE_CURRENT_SOURCE_DIR}/metal/ops/rmsnorm_fw/device/rmsnorm_fw_program_factory.cpp
+    # CrossEntropy
+    ${CMAKE_CURRENT_SOURCE_DIR}/metal/ops/cross_entropy_fw/cross_entropy_fw.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/metal/ops/cross_entropy_fw/cross_entropy_fw.hpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/metal/ops/cross_entropy_fw/device/cross_entropy_fw_device_operation_types.hpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/metal/ops/cross_entropy_fw/device/cross_entropy_fw_device_operation.hpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/metal/ops/cross_entropy_fw/device/cross_entropy_fw_device_operation.cpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/metal/ops/cross_entropy_fw/device/cross_entropy_fw_program_factory.hpp
+    ${CMAKE_CURRENT_SOURCE_DIR}/metal/ops/cross_entropy_fw/device/cross_entropy_fw_program_factory.cpp
 )
 
 list(APPEND SOURCES ${METAL_OPS_FILES})
diff --git a/tt-train/sources/ttml/metal/operations.hpp b/tt-train/sources/ttml/metal/operations.hpp
index 11cd5cab147c..776e56b1f3df 100644
--- a/tt-train/sources/ttml/metal/operations.hpp
+++ b/tt-train/sources/ttml/metal/operations.hpp
@@ -2,6 +2,9 @@
 //
 // SPDX-License-Identifier: Apache-2.0
 
+#pragma once
+
+#include "ops/cross_entropy_fw/cross_entropy_fw.hpp"
 #include "ops/rmsnorm_fw/rmsnorm_fw.hpp"
 
 namespace ttml::metal {
@@ -9,4 +12,8 @@ namespace ttml::metal {
 constexpr auto rmsnorm_fw =
     ttnn::register_operation<"ttml::metal::rmsnorm_fw", ttml::metal::ops::rmsnorm_fw::RMSNormForwardOperation>();
 
+constexpr auto cross_entropy_fw = ttnn::register_operation<
+    "ttml::metal::cross_entropy_fw",
+    ttml::metal::ops::cross_entropy_fw::CrossEntropyForwardOperation>();
+
 }  // namespace ttml::metal
diff --git a/tt-train/sources/ttml/metal/ops/cross_entropy_fw/cross_entropy_fw.cpp b/tt-train/sources/ttml/metal/ops/cross_entropy_fw/cross_entropy_fw.cpp
new file mode 100644
index 000000000000..4bf2492fbf1b
--- /dev/null
+++ b/tt-train/sources/ttml/metal/ops/cross_entropy_fw/cross_entropy_fw.cpp
@@ -0,0 +1,15 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include "cross_entropy_fw.hpp"
+
+#include "device/cross_entropy_fw_device_operation.hpp"
+
+namespace ttml::metal::ops::cross_entropy_fw {
+
+ttnn::Tensor CrossEntropyForwardOperation::invoke(const ttnn::Tensor& input_tensor, const ttnn::Tensor& target_tensor) {
+    auto result = ttnn::prim::ttml_cross_entropy_fw(input_tensor, target_tensor);
+    return result;
+}
+}  // namespace ttml::metal::ops::cross_entropy_fw
diff --git a/tt-train/sources/ttml/metal/ops/cross_entropy_fw/cross_entropy_fw.hpp b/tt-train/sources/ttml/metal/ops/cross_entropy_fw/cross_entropy_fw.hpp
new file mode 100644
index 000000000000..d01412b191ee
--- /dev/null
+++ b/tt-train/sources/ttml/metal/ops/cross_entropy_fw/cross_entropy_fw.hpp
@@ -0,0 +1,17 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#pragma once
+
+#include "metal/ttnn_all_includes.hpp"
+
+namespace ttml::metal::ops::cross_entropy_fw {
+
+struct CrossEntropyForwardOperation {
+    static ttnn::Tensor invoke(
+        const ttnn::Tensor& input,  // logits : model output (N, 1, H, W)
+        const ttnn::Tensor& target  // target : ground truth (N, H)
+    );
+};
+}  // namespace ttml::metal::ops::cross_entropy_fw
diff --git a/tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/cross_entropy_fw_device_operation.cpp b/tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/cross_entropy_fw_device_operation.cpp
new file mode 100644
index 000000000000..9d97c2cd6ae4
--- /dev/null
+++ b/tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/cross_entropy_fw_device_operation.cpp
@@ -0,0 +1,128 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include "cross_entropy_fw_device_operation.hpp"
+
+#include "cross_entropy_fw_program_factory.hpp"
+
+namespace ttml::metal::ops::cross_entropy_fw::device {
+
+CrossEntropyForwardDeviceOperation::program_factory_t CrossEntropyForwardDeviceOperation::select_program_factory(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    return CrossEntropyForwardProgramFactory{};
+}
+
+void CrossEntropyForwardDeviceOperation::validate_on_program_cache_hit(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    validate_on_program_cache_miss(args, tensor_args);
+}
+
+void CrossEntropyForwardDeviceOperation::validate_on_program_cache_miss(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    auto check_tensor = [](const ttnn::Tensor& tensor,
+                           const std::string& name,
+                           const tt::tt_metal::Layout required_layout,
+                           const tt::tt_metal::DataType required_dtype) {
+        TT_FATAL(
+            tensor.device()->arch() == tt::ARCH::WORMHOLE_B0,
+            "CrossEntropyForward operation is only supported on Wormhole. Device arch: {}. Tensor name: {}",
+            magic_enum::enum_name(tensor.device()->arch()),
+            name);
+
+        TT_FATAL(
+            tensor.storage_type() == tt::tt_metal::StorageType::DEVICE,
+            "CrossEntropyForward operation requires '{}' to be on DEVICE. Got storage type: '{}'",
+            name,
+            magic_enum::enum_name(tensor.storage_type()));
+
+        TT_FATAL(tensor.buffer() != nullptr, "Tensor '{}' must be allocated on device (buffer is null).", name);
+
+        TT_FATAL(
+            tensor.get_layout() == required_layout,
+            "Tensor '{}' must have layout '{}', but got '{}'",
+            name,
+            magic_enum::enum_name(required_layout),
+            magic_enum::enum_name(tensor.get_layout()));
+
+        TT_FATAL(
+            tensor.get_dtype() == required_dtype,
+            "Tensor '{}' must have data type '{}', but got '{}'",
+            name,
+            magic_enum::enum_name(required_dtype),
+            magic_enum::enum_name(tensor.get_dtype()));
+
+        TT_FATAL(
+            tensor.memory_config().memory_layout == ttnn::TensorMemoryLayout::INTERLEAVED,
+            "Tensor '{}' must use INTERLEAVED memory layout, but got '{}'",
+            name,
+            magic_enum::enum_name(tensor.memory_config().memory_layout));
+    };
+
+    const auto& input_tensor = tensor_args.input;
+    const auto& target_tensor = tensor_args.target;
+    const auto& preallocated_output_tensor = tensor_args.preallocated_output;
+    check_tensor(input_tensor, "Input", tt::tt_metal::Layout::TILE, tt::tt_metal::DataType::BFLOAT16);
+    check_tensor(target_tensor, "Target", tt::tt_metal::Layout::ROW_MAJOR, tt::tt_metal::DataType::UINT32);
+    if (preallocated_output_tensor.has_value()) {
+        check_tensor(
+            preallocated_output_tensor.value(),
+            "Preallocated Output",
+            tt::tt_metal::Layout::TILE,
+            tt::tt_metal::DataType::BFLOAT16);
+    }
+}
+
+CrossEntropyForwardDeviceOperation::spec_return_value_t CrossEntropyForwardDeviceOperation::compute_output_specs(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    if (tensor_args.preallocated_output.has_value()) {
+        return tensor_args.preallocated_output->get_tensor_spec();
+    }
+    auto input_logical_shape = tensor_args.input.get_logical_shape();
+    input_logical_shape[-1] = 1U;
+    return ttnn::TensorSpec(
+        ttnn::Shape(input_logical_shape),
+        tt::tt_metal::TensorLayout(
+            tensor_args.input.get_dtype(), tt::tt_metal::Layout::TILE, tensor_args.input.memory_config()));
+}
+
+CrossEntropyForwardDeviceOperation::tensor_return_value_t CrossEntropyForwardDeviceOperation::create_output_tensors(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    tensor_return_value_t output_tensor;
+
+    spec_return_value_t output_specs = compute_output_specs(args, tensor_args);
+
+    if (tensor_args.preallocated_output.has_value()) {
+        output_tensor = tensor_args.preallocated_output.value();
+    } else {
+        output_tensor = create_device_tensor(output_specs, tensor_args.input.device());
+    }
+
+    return output_tensor;
+}
+
+tt::stl::hash::hash_t CrossEntropyForwardDeviceOperation::compute_program_hash(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    const auto& input_tensor = tensor_args.input;
+    const auto& input_logical_shape = input_tensor.get_logical_shape();
+    auto program_factory = select_program_factory(args, tensor_args);
+    auto hash = tt::tt_metal::operation::hash_operation<CrossEntropyForwardDeviceOperation>(
+        args, program_factory.index(), input_tensor.dtype(), input_logical_shape);
+
+    return hash;
+}
+
+std::tuple<operation_attributes_t, tensor_args_t> CrossEntropyForwardDeviceOperation::invoke(
+    const ttnn::Tensor& input_tensor,
+    const ttnn::Tensor& target_tensor,
+    const std::optional<ttnn::Tensor>& preallocated_output) {
+    return {
+        operation_attributes_t{},
+        tensor_args_t{
+            .input = input_tensor,
+            .target = target_tensor,
+            .preallocated_output = preallocated_output,
+        }};
+}
+
+}  // namespace ttml::metal::ops::cross_entropy_fw::device
diff --git a/tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/cross_entropy_fw_device_operation.hpp b/tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/cross_entropy_fw_device_operation.hpp
new file mode 100644
index 000000000000..72bf4fd184c1
--- /dev/null
+++ b/tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/cross_entropy_fw_device_operation.hpp
@@ -0,0 +1,47 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#pragma once
+
+#include <optional>
+
+#include "cross_entropy_fw_device_operation_types.hpp"
+#include "cross_entropy_fw_program_factory.hpp"
+#include "metal/ttnn_all_includes.hpp"
+
+namespace ttml::metal::ops::cross_entropy_fw::device {
+
+struct CrossEntropyForwardDeviceOperation {
+    using operation_attributes_t = operation_attributes_t;
+    using tensor_args_t = tensor_args_t;
+    using spec_return_value_t = spec_return_value_t;
+    using tensor_return_value_t = tensor_return_value_t;
+    using program_factory_t = std::variant<CrossEntropyForwardProgramFactory>;
+
+    static program_factory_t select_program_factory(const operation_attributes_t&, const tensor_args_t&);
+
+    static void validate_on_program_cache_hit(const operation_attributes_t&, const tensor_args_t&);
+    static void validate_on_program_cache_miss(const operation_attributes_t&, const tensor_args_t&);
+
+    static spec_return_value_t compute_output_specs(const operation_attributes_t&, const tensor_args_t&);
+
+    static tensor_return_value_t create_output_tensors(
+        const operation_attributes_t& operation_attributes, const tensor_args_t&);
+
+    static tt::stl::hash::hash_t compute_program_hash(const operation_attributes_t&, const tensor_args_t&);
+
+    static std::tuple<operation_attributes_t, tensor_args_t> invoke(
+        const ttnn::Tensor& input_tensor,
+        const ttnn::Tensor& target_tensor,
+        const std::optional<ttnn::Tensor>& preallocated_output = std::nullopt);
+};
+
+}  // namespace ttml::metal::ops::cross_entropy_fw::device
+
+namespace ttnn::prim {
+
+constexpr auto ttml_cross_entropy_fw = ttnn::register_operation<
+    "ttnn::prim::ttml_cross_entropy_fw",
+    ttml::metal::ops::cross_entropy_fw::device::CrossEntropyForwardDeviceOperation>();
+}  // namespace ttnn::prim
diff --git a/tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/cross_entropy_fw_device_operation_types.hpp b/tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/cross_entropy_fw_device_operation_types.hpp
new file mode 100644
index 000000000000..a92c33e4f4c3
--- /dev/null
+++ b/tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/cross_entropy_fw_device_operation_types.hpp
@@ -0,0 +1,23 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#pragma once
+
+#include "metal/ttnn_all_includes.hpp"
+
+namespace ttml::metal::ops::cross_entropy_fw::device {
+
+struct operation_attributes_t {};
+
+struct tensor_args_t {
+    const ttnn::Tensor& input;
+    const ttnn::Tensor& target;
+
+    std::optional<ttnn::Tensor> preallocated_output;
+};
+
+using tensor_return_value_t = ttnn::Tensor;  // return loss: tensor with shape (N, 1, H, 1)
+using spec_return_value_t = ttnn::TensorSpec;
+
+}  // namespace ttml::metal::ops::cross_entropy_fw::device
diff --git a/tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/cross_entropy_fw_program_factory.cpp b/tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/cross_entropy_fw_program_factory.cpp
new file mode 100644
index 000000000000..bb69076db839
--- /dev/null
+++ b/tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/cross_entropy_fw_program_factory.cpp
@@ -0,0 +1,504 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include "cross_entropy_fw_program_factory.hpp"
+
+#include <bit>
+#include <cstdint>
+#include <tt-metalium/buffer.hpp>
+
+#include "cross_entropy_fw_device_operation_types.hpp"
+
+namespace {
+
+constexpr auto kReaderKernelPath =
+    "tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/kernels/dataflow/"
+    "reader_cross_entropy_fw_interleaved_start_id.cpp";
+
+constexpr auto kWriterKernelPath =
+    "tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/kernels/dataflow/"
+    "writer_cross_entropy_fw_interleaved_start_id.cpp";
+
+constexpr auto kComputeKernelPath =
+    "tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/kernels/compute/cross_entropy_fw_kernel.cpp";
+
+// reader runtime args
+constexpr uint32_t kInputBufferIdx = 0;
+constexpr uint32_t kTargetBufferIdx = 1U;
+// writer runtime args
+constexpr uint32_t kOutputBufferIdx = 0;
+
+constexpr auto kInputCbIndex = tt::CBIndex::c_0;
+constexpr auto kTargetCbIndex = tt::CBIndex::c_1;
+constexpr auto kMaskCbIndex = tt::CBIndex::c_2;
+constexpr auto kMaxMaskCbIndex = tt::CBIndex::c_3;
+constexpr auto kScalerCbIndex = tt::CBIndex::c_4;  // used to reduction
+constexpr auto kInputTileCbIndex = tt::CBIndex::c_5;
+constexpr auto kTargetLogitsCbIndex = tt::CBIndex::c_6;
+constexpr auto kMaxValueBeforeReductionCbIndex = tt::CBIndex::c_7;
+constexpr auto kMaxValueAfterReductionCbIndex = tt::CBIndex::c_8;
+constexpr auto kExpSumBeforeReductionCbIndex = tt::CBIndex::c_9;
+constexpr auto KExpSumAfterReductionCbIndex = tt::CBIndex::c_10;
+constexpr auto kOutputCbIndex = tt::CBIndex::c_11;
+
+constexpr uint32_t kNumTargetIndexesTiles = 1U;
+constexpr uint32_t kNumMaskTiles = 1U;
+constexpr uint32_t kNumMaxValueTiles = 1U;
+constexpr uint32_t kNumInputTilesByIndx = 1U;
+constexpr uint32_t kMaxValueBeforeReductionTiles = 2U;
+constexpr uint32_t kNumTargetLogitsTiles = 2U;
+constexpr uint32_t kNumMaxValueAfterReductionTiles = 2U;
+constexpr uint32_t kNumExpSumBeforeReductionTiles = 2U;
+constexpr uint32_t kNumExpSumAfterReductionTiles = 2U;
+constexpr uint32_t kNumScalerTiles = 1U;  // used it to reduction
+constexpr uint32_t kNumOutputTiles = 1U;
+
+constexpr uint32_t kPageElementsNumber = 32U;
+
+const std::string kMaskWDefineKey = "DO_MASK_W";
+const std::string kEverythingFitsInL1DefineKey = "EVERYTHING_FITS_IN_L1";
+
+uint32_t get_block_size(uint32_t num_inner) {
+    const uint32_t max_block_size = 4U;  // 4 is the maximum block size for enabled fp32 dest acc
+    for (uint32_t block_size = max_block_size; block_size > 1; block_size--) {
+        if (num_inner % block_size == 0) {  // if num_inner is divisible by block_size - choose this block_size
+            return block_size;
+        }
+    }
+    return 1U;
+}
+
+}  // namespace
+
+namespace ttml::metal::ops::cross_entropy_fw::device {
+
+/**
+ *   Helper struct to hold references to all kernels we create,
+ *        used during runtime argument setup.
+ */
+struct CrossEntropyForwardKernels {
+    tt::tt_metal::KernelHandle reader;
+    tt::tt_metal::KernelHandle writer;
+    tt::tt_metal::KernelHandle compute_group_1;
+    tt::tt_metal::KernelHandle compute_group_2;
+};
+
+/**
+ *   Create and configure a circular buffer, returning both the configuration and the handle.
+ */
+tt::tt_metal::CBHandle create_circular_buffer(
+    tt::tt_metal::Program& program,
+    const tt::tt_metal::CoreRangeSet& core_ranges,
+    uint32_t cb_index,
+    tt::DataFormat data_format,
+    uint32_t single_tile_size,
+    uint32_t num_tiles) {
+    tt::tt_metal::CircularBufferConfig cb_config =
+        tt::tt_metal::CircularBufferConfig(num_tiles * single_tile_size, {{cb_index, data_format}})
+            .set_page_size(cb_index, single_tile_size);
+
+    auto cb_handle = CreateCircularBuffer(program, core_ranges, cb_config);
+    return cb_handle;
+}
+
+/**
+ *   Create a reader kernel with the given compile-time arguments.
+ */
+tt::tt_metal::KernelHandle create_reader_kernel(
+    tt::tt_metal::Program& program,
+    const tt::tt_metal::CoreRangeSet& core_ranges,
+    const std::vector<uint32_t>& compile_time_args,
+    const std::map<std::string, std::string>& defines,
+    const std::string& kernel_path) {
+    return tt::tt_metal::CreateKernel(
+        program, kernel_path, core_ranges, tt::tt_metal::ReaderDataMovementConfig(compile_time_args, defines));
+}
+
+/**
+ *   Create a writer kernel with the given compile-time arguments.
+ */
+tt::tt_metal::KernelHandle create_writer_kernel(
+    tt::tt_metal::Program& program,
+    const tt::tt_metal::CoreRangeSet& core_ranges,
+    const std::vector<uint32_t>& compile_time_args,
+    const std::map<std::string, std::string>& defines,
+    const std::string& kernel_path) {
+    return tt::tt_metal::CreateKernel(
+        program, kernel_path, core_ranges, tt::tt_metal::WriterDataMovementConfig(compile_time_args, defines));
+}
+
+/**
+ * Create a compute kernel with the given compile-time arguments.
+ */
+tt::tt_metal::KernelHandle create_compute_kernel(
+    tt::tt_metal::Program& program,
+    const tt::tt_metal::CoreRangeSet& core_ranges,
+    const std::vector<uint32_t>& compile_time_args,
+    const std::map<std::string, std::string>& defines,
+    const std::string& kernel_path) {
+    return tt::tt_metal::CreateKernel(
+        program,
+        kernel_path,
+        core_ranges,
+        tt::tt_metal::ComputeConfig{
+            .math_fidelity = MathFidelity::HiFi4,
+            .fp32_dest_acc_en = true,
+            .math_approx_mode = false,
+            .compile_args = compile_time_args,
+            .defines = defines});
+}
+
+/**
+ * Set up the runtime arguments for the 4 relevant kernels (reader, writer, compute G1, compute G2)
+ *        for each core in the grid.
+ */
+void assign_per_core_runtime_args(
+    tt::tt_metal::Program& program,
+    const CrossEntropyForwardKernels& kernels,
+    const tt::tt_metal::Buffer* input_buffer,
+    const tt::tt_metal::Buffer* target_buffer,
+    const tt::tt_metal::Buffer* output_buffer,
+
+    uint32_t num_cores,
+    uint32_t num_cores_y,
+    uint32_t num_rows_per_core_group_1,
+    uint32_t num_rows_per_core_group_2,
+    const tt::tt_metal::CoreRangeSet& core_group_1,
+    const tt::tt_metal::CoreRangeSet& core_group_2) {
+    for (uint32_t i = 0, num_rows_written = 0; i < num_cores; i++) {
+        CoreCoord core = {i / num_cores_y, i % num_cores_y};
+
+        // Determine how many rows this core will process
+        uint32_t num_rows_per_core = 0;
+        if (core_group_1.contains(core)) {
+            num_rows_per_core = num_rows_per_core_group_1;
+        } else if (core_group_2.contains(core)) {
+            num_rows_per_core = num_rows_per_core_group_2;
+        } else {
+            TT_FATAL(false, "Core not in specified core ranges");
+        }
+
+        // Reader kernel: (input_addr, gamma_addr, number_of_rows, offset_in_rows)
+        SetRuntimeArgs(
+            program,
+            kernels.reader,
+            core,
+            {input_buffer->address(), target_buffer->address(), num_rows_per_core, num_rows_written});
+
+        // Writer kernel: (dst_addr, dst_rms_addr number_of_rows, offset_in_rows)
+        SetRuntimeArgs(
+            program,
+            kernels.writer,
+            core,
+            {output_buffer->address(), /*temporary*/ num_rows_per_core, num_rows_written});
+
+        num_rows_written += num_rows_per_core;
+    }
+}
+
+CrossEntropyForwardProgramFactory::cached_program_t CrossEntropyForwardProgramFactory::create(
+    const operation_attributes_t& operation_attributes,
+    const tensor_args_t& tensor_args,
+    tensor_return_value_t& output) {
+    // -------------------------------------------------------------------------
+    // 1) Setup device, data formats, tile sizes, and compute split
+    // -------------------------------------------------------------------------
+    const auto& input = tensor_args.input;
+    const auto& target = tensor_args.target;
+
+    auto* device = input.device();
+
+    tt::tt_metal::Program program{};
+
+    tt::DataFormat input_data_format = datatype_to_dataformat_converter(input.get_dtype());
+    TT_FATAL(input_data_format == tt::DataFormat::Float16_b, "Input data format must be Float16_b");
+
+    uint32_t bfloat16_single_tile_size_bytes = tt::tt_metal::detail::TileSize(tt::DataFormat::Float16_b);
+    uint32_t float32_single_tile_size_bytes = tt::tt_metal::detail::TileSize(tt::DataFormat::Float32);
+
+    auto padded_tensor_shape = input.get_padded_shape();
+    auto padded_tensor_volume = input.volume();
+
+    TT_FATAL(
+        padded_tensor_volume % tt::constants::TILE_HW == 0, "Padded input tensor volume must be divisible by TILE_HW");
+    TT_FATAL(padded_tensor_shape.rank() == 4U, "Input tensor must be 4D");
+
+    //
+    uint32_t Wt = padded_tensor_shape[-1] / tt::constants::TILE_WIDTH;  // <- number of tiles in inner dimension
+    uint32_t Ht = padded_tensor_shape[-2] / tt::constants::TILE_HEIGHT;
+    uint32_t NC = padded_tensor_shape[0] * padded_tensor_shape[1];
+    uint32_t total_rows_to_process = NC * Ht;
+
+    // get size of target indexes inner dimension
+    uint32_t target_indexes_inner_dim_size = target.get_logical_shape()[-1] * target.element_size();
+    // read target indexes by pages(32 indexes in page)
+    uint32_t uint32_read_page_size = tt::datum_size(tt::DataFormat::UInt32) * kPageElementsNumber;
+
+    // get number of free cores
+    auto compute_with_storage_grid_size = device->compute_with_storage_grid_size();
+    uint32_t num_cores_x = compute_with_storage_grid_size.x;
+    uint32_t num_cores_y = compute_with_storage_grid_size.y;
+
+    // get the number of inner dimension
+    uint32_t num_inner = input.get_logical_shape()[-1];  // (N, 1, C, H)
+
+    // mask_w - this mask used to avoid calculation of extra data(data which will be added to create full tile 32x32)??
+    uint32_t mask_w = num_inner % tt::constants::TILE_WIDTH;  // width index of first trash value in tile
+
+    // compile arguments
+    uint32_t block_size = get_block_size(Wt);
+
+    auto [num_cores, all_cores, core_group_1, core_group_2, num_rows_per_core_group_1, num_rows_per_core_group_2] =
+        tt::tt_metal::split_work_to_cores(compute_with_storage_grid_size, total_rows_to_process);
+
+    // -------------------------------------------------------------------------
+    // 2) Create and configure circular buffers
+    // -------------------------------------------------------------------------
+
+    uint32_t twice_block_size = 2U * block_size;
+
+    const uint32_t available_L1_in_bytes =
+        device->l1_size_per_core() - device->allocator()->get_base_allocator_addr(tt::tt_metal::HalMemType::L1);
+
+    const uint64_t mask_scaler_maxmask_memory =
+        (kNumMaskTiles + kNumScalerTiles + kNumMaskTiles) * bfloat16_single_tile_size_bytes;
+    const uint64_t output_memory = kNumOutputTiles * bfloat16_single_tile_size_bytes;
+    const uint64_t max_value_memory =
+        (kNumMaxValueAfterReductionTiles + kMaxValueBeforeReductionTiles) * bfloat16_single_tile_size_bytes;
+    const uint64_t exp_sum_memory =
+        (kNumExpSumBeforeReductionTiles + kNumExpSumAfterReductionTiles) * float32_single_tile_size_bytes;
+    const uint64_t input_memory =
+        Wt * bfloat16_single_tile_size_bytes + 2U * bfloat16_single_tile_size_bytes + uint32_read_page_size;
+
+    // Total L1 memory required
+    const uint64_t required_L1_in_bytes =
+        input_memory + mask_scaler_maxmask_memory + max_value_memory + exp_sum_memory + output_memory;
+    // Is everything fits in L1
+    const bool everything_fits_in_l1 = required_L1_in_bytes <= available_L1_in_bytes;
+
+    const uint32_t num_input_tiles = (everything_fits_in_l1) ? Wt : twice_block_size;
+
+    auto data_format = input_data_format;  // tt::DataFormat::Float16_b
+    auto precise_data_format = tt::DataFormat::Float32;
+    auto target_indexes_data_format = tt::DataFormat::UInt32;
+
+    auto cb_input = create_circular_buffer(
+        program, all_cores, kInputCbIndex, data_format, bfloat16_single_tile_size_bytes, num_input_tiles);
+
+    auto cb_target = create_circular_buffer(
+        program, all_cores, kTargetCbIndex, target_indexes_data_format, uint32_read_page_size, kNumTargetIndexesTiles);
+
+    auto cb_mask = create_circular_buffer(
+        program, all_cores, kMaskCbIndex, data_format, bfloat16_single_tile_size_bytes, kNumMaskTiles);
+
+    auto cb_max_mask = create_circular_buffer(
+        program, all_cores, kMaxMaskCbIndex, data_format, bfloat16_single_tile_size_bytes, kNumMaskTiles);
+
+    auto cb_scaler = create_circular_buffer(
+        program, all_cores, kScalerCbIndex, data_format, bfloat16_single_tile_size_bytes, kNumScalerTiles);
+
+    auto cb_input_tile_by_idx = create_circular_buffer(
+        program, all_cores, kInputTileCbIndex, data_format, bfloat16_single_tile_size_bytes, kNumInputTilesByIndx);
+
+    auto cb_target_inputs = create_circular_buffer(
+        program, all_cores, kTargetLogitsCbIndex, data_format, bfloat16_single_tile_size_bytes, kNumTargetLogitsTiles);
+
+    auto cb_max_value_before_reduction = create_circular_buffer(
+        program,
+        all_cores,
+        kMaxValueBeforeReductionCbIndex,
+        data_format,
+        bfloat16_single_tile_size_bytes,
+        kMaxValueBeforeReductionTiles);
+
+    auto cb_max_value_after_reduction = create_circular_buffer(
+        program,
+        all_cores,
+        kMaxValueAfterReductionCbIndex,
+        data_format,
+        bfloat16_single_tile_size_bytes,
+        kNumMaxValueAfterReductionTiles);
+
+    auto cb_exp_sum_before_reduction = create_circular_buffer(
+        program,
+        all_cores,
+        kExpSumBeforeReductionCbIndex,
+        data_format,
+        bfloat16_single_tile_size_bytes,
+        kNumExpSumBeforeReductionTiles);
+
+    auto cb_exp_sum_after_refuction = create_circular_buffer(
+        program,
+        all_cores,
+        KExpSumAfterReductionCbIndex,
+        data_format,
+        bfloat16_single_tile_size_bytes,
+        kNumExpSumAfterReductionTiles);
+
+    auto cb_output = create_circular_buffer(
+        program,
+        all_cores,
+        kOutputCbIndex,
+        data_format,
+        bfloat16_single_tile_size_bytes,
+        kNumOutputTiles);  // create twice block_size variable it looks like I need 1 tile here
+
+    // -------------------------------------------------------------------------
+    // 3) Create reader/writer kernels
+    // -------------------------------------------------------------------------
+
+    auto* input_buffer = input.buffer();
+    TT_FATAL(
+        input_buffer->buffer_type() == ttnn::BufferType::DRAM,
+        "Input buffer must be in DRAM. Input buffer of type {}",
+        magic_enum::enum_name(input_buffer->buffer_type()));
+
+    auto* target_buffer = target.buffer();
+    TT_FATAL(
+        target_buffer->buffer_type() == ttnn::BufferType::DRAM,
+        "Target buffer must be in DRAM. Target buffer of type {}",
+        magic_enum::enum_name(target_buffer->buffer_type()));
+
+    auto* output_buffer = output.buffer();
+    TT_FATAL(
+        output_buffer->buffer_type() == ttnn::BufferType::DRAM,
+        "Output buffer must be in DRAM. Output buffer of type {}",
+        magic_enum::enum_name(output_buffer->buffer_type()));
+
+    // configure defines
+    std::map<std::string, std::string> defines;
+    // define whether I need mask or not
+    if (mask_w != 0) {
+        defines[kMaskWDefineKey] = "1";
+    }
+
+    // define whether everything fits in L1 or not
+    if (everything_fits_in_l1) {
+        defines[kEverythingFitsInL1DefineKey] = "1";
+    }
+
+    // setup defines for reduce
+    // Compute kernel does not compile without these defines
+    // LLK reduction uses define values as default template parameters
+    defines["REDUCE_OP"] = "PoolType::SUM";
+    defines["REDUCE_DIM"] = "ReduceDim::REDUCE_ROW";
+
+    CrossEntropyForwardKernels kernels;
+    kernels.reader = create_reader_kernel(
+        program,
+        all_cores,
+        /* reader_compile_args */
+        {block_size, Wt, mask_w, target_indexes_inner_dim_size, Ht, uint32_read_page_size},
+        defines,
+        kReaderKernelPath);
+
+    kernels.writer = create_writer_kernel(
+        program, all_cores, /* writer_compile_args */ {block_size, Wt}, defines, kWriterKernelPath);
+
+    // -------------------------------------------------------------------------
+    // 4) Create compute kernels for cross_entropy_fw
+    // -------------------------------------------------------------------------
+
+    // Group 1 compile-time arguments
+    std::vector<uint32_t> compute_group_1_args = {
+        num_rows_per_core_group_1,  // per_core_block_cnt
+        block_size,                 // per_core_block_size
+        Wt                          // num_inner / TILE_W
+    };
+
+    kernels.compute_group_1 =
+        create_compute_kernel(program, core_group_1, compute_group_1_args, defines, kComputeKernelPath);
+
+    // Group 2 (if present) compile-time arguments
+    std::vector<uint32_t> compute_group_2_args = {
+        num_rows_per_core_group_2,  // per_core_block_cnt
+        block_size,                 // per_core_block_size
+        Wt                          // num_inner / TILE_W
+    };
+
+    kernels.compute_group_2 =
+        create_compute_kernel(program, core_group_2, compute_group_2_args, defines, kComputeKernelPath);
+
+    // -------------------------------------------------------------------------
+    // 5) Assign runtime args for each core
+    // -------------------------------------------------------------------------
+
+    assign_per_core_runtime_args(
+        program,
+        kernels,
+        input_buffer,
+        target_buffer,
+        output_buffer,
+        num_cores,
+        num_cores_y,
+        num_rows_per_core_group_1,
+        num_rows_per_core_group_2,
+        core_group_1,
+        core_group_2);
+
+    // -------------------------------------------------------------------------
+    // 6) Return the fully configured program & relevant shared variables
+    // -------------------------------------------------------------------------
+
+    return cached_program_t{
+        std::move(program),
+        {/* rmsnorm_fw_reader_kernel_id  = */ kernels.reader,
+         /* rmsnorm_fw_writer_kernel_id  = */ kernels.writer,
+         /* rmsnorm_fw_kernel_group_1_id = */ kernels.compute_group_1,
+         /* rmsnorm_fw_kernel_group_2_id = */ kernels.compute_group_2,
+         /* core_group_1              = */ core_group_1,
+         /* core_group_2              = */ core_group_2,
+         /* num_cores                 = */ num_cores,
+         /* num_cores_y               = */ num_cores_y}};
+}
+
+void CrossEntropyForwardProgramFactory::override_runtime_arguments(
+    cached_program_t& cached_program,
+    const operation_attributes_t& operation_attributes,
+    const tensor_args_t& tensor_args,
+    tensor_return_value_t& tensor_return_value) {
+    auto& program = cached_program.program;
+    auto& shared_variables = cached_program.shared_variables;
+    auto& cross_entropy_fw_reader_kernel_id = shared_variables.reader_kernel_id;
+    auto& cross_entropy_fw_writer_kernel_id = shared_variables.writer_kernel_id;
+    auto& cross_entropy_fw_kernel_group_1_id = shared_variables.compute_kernel_group_1_id;
+    auto& cross_entropy_fw_kernel_group_2_id = shared_variables.compute_kernel_group_2_id;
+    auto& core_group_1 = shared_variables.core_group_1;
+    auto& core_group_2 = shared_variables.core_group_2;
+
+    uint32_t num_cores = shared_variables.num_cores;
+    uint32_t num_cores_y = shared_variables.num_cores_y;
+
+    auto* input_buffer = tensor_args.input.buffer();
+    auto* target_buffer = tensor_args.target.buffer();
+    auto* output_buffer = tensor_return_value.buffer();
+
+    // Only address arguments need updating here; tile counts remain the same as in create().
+    auto& reader_runtime_args = GetRuntimeArgs(program, cross_entropy_fw_reader_kernel_id);
+    auto& writer_runtime_args = GetRuntimeArgs(program, cross_entropy_fw_writer_kernel_id);
+    auto& group_1_runtime_args = GetRuntimeArgs(program, cross_entropy_fw_kernel_group_1_id);
+    // we need to initialize it with something, but if group 2 is  empty it will be used in the loop
+    auto& group_2_runtime_args = core_group_2.ranges().empty()
+                                     ? group_1_runtime_args
+                                     : GetRuntimeArgs(program, cross_entropy_fw_kernel_group_2_id);
+
+    for (uint32_t i = 0; i < num_cores; i++) {
+        CoreCoord core = {i / num_cores_y, i % num_cores_y};
+
+        // Update input buffers for the reader kernel
+        {
+            auto& runtime_args = reader_runtime_args[core.x][core.y];
+            runtime_args[kInputBufferIdx] = input_buffer->address();
+            runtime_args[kTargetBufferIdx] = target_buffer->address();
+        }
+
+        // Update output buffers for the writer kernel
+        {
+            auto& runtime_args = writer_runtime_args[core.x][core.y];
+            runtime_args[kOutputBufferIdx] = output_buffer->address();
+        }
+    }
+}
+
+}  // namespace ttml::metal::ops::cross_entropy_fw::device
diff --git a/tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/cross_entropy_fw_program_factory.hpp b/tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/cross_entropy_fw_program_factory.hpp
new file mode 100644
index 000000000000..f51fca08c190
--- /dev/null
+++ b/tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/cross_entropy_fw_program_factory.hpp
@@ -0,0 +1,37 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#pragma once
+
+#include "cross_entropy_fw_device_operation_types.hpp"
+#include "metal/ttnn_all_includes.hpp"
+
+namespace ttml::metal::ops::cross_entropy_fw::device {
+
+struct CrossEntropyForwardProgramFactory {
+    struct shared_variables_t {
+        tt::tt_metal::KernelHandle reader_kernel_id;
+        tt::tt_metal::KernelHandle writer_kernel_id;
+        tt::tt_metal::KernelHandle compute_kernel_group_1_id;
+        tt::tt_metal::KernelHandle compute_kernel_group_2_id;
+        CoreRangeSet core_group_1;
+        CoreRangeSet core_group_2;
+        uint32_t num_cores{};
+        uint32_t num_cores_y{};
+    };
+    using cached_program_t = ttnn::device_operation::CachedProgram<shared_variables_t>;
+
+    static cached_program_t create(
+        const operation_attributes_t& operation_attributes,
+        const tensor_args_t& tensor_args,
+        tensor_return_value_t& tensor_return_value);
+
+    static void override_runtime_arguments(
+        cached_program_t& cached_program,
+        const operation_attributes_t& operation_attributes,
+        const tensor_args_t& tensor_args,
+        tensor_return_value_t& tensor_return_value);
+};
+
+}  // namespace ttml::metal::ops::cross_entropy_fw::device
diff --git a/tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/kernels/compute/cross_entropy_fw_kernel.cpp b/tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/kernels/compute/cross_entropy_fw_kernel.cpp
new file mode 100644
index 000000000000..cb018f13b0b4
--- /dev/null
+++ b/tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/kernels/compute/cross_entropy_fw_kernel.cpp
@@ -0,0 +1,406 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include <compute_kernel_api/cb_api.h>
+#include <compute_kernel_api/pack.h>
+#include <compute_kernel_api/reconfig_data_format.h>
+#include <compute_kernel_api/reg_api.h>
+#include <debug/dprint.h>
+#include <hostdevcommon/kernel_structs.h>
+#include <tensix.h>
+
+#include <cstdint>
+
+#include "compute_kernel_api.h"
+#include "compute_kernel_api/bcast.h"
+#include "compute_kernel_api/common.h"
+#include "compute_kernel_api/eltwise_binary.h"
+#include "compute_kernel_api/eltwise_binary_sfpu.h"
+#include "compute_kernel_api/eltwise_unary/eltwise_unary.h"
+#include "compute_kernel_api/eltwise_unary/exp.h"
+#include "compute_kernel_api/eltwise_unary/negative.h"
+#include "compute_kernel_api/eltwise_unary/recip.h"
+#include "compute_kernel_api/eltwise_unary/sfpu_split_includes.h"
+#include "compute_kernel_api/eltwise_unary/sqrt.h"
+#include "compute_kernel_api/mask.h"
+#include "compute_kernel_api/reduce.h"
+#include "compute_kernel_api/tile_move_copy.h"
+
+namespace NAMESPACE {
+
+constexpr uint32_t num_rows_per_core = get_compile_time_arg_val(0);  // rows to process in this kernel
+constexpr uint32_t block_size = get_compile_time_arg_val(1);         // size of block
+constexpr uint32_t Wt = get_compile_time_arg_val(2);                 // number of tiles in inner dimension
+
+constexpr auto cb_input = tt::CBIndex::c_0;
+constexpr auto cb_mask = tt::CBIndex::c_2;
+constexpr auto cb_max_mask = tt::CBIndex::c_3;
+constexpr auto cb_scaler = tt::CBIndex::c_4;  // used to reduction
+constexpr auto cb_target_logits = tt::CBIndex::c_6;
+constexpr auto cb_max_value_before_reduction = tt::CBIndex::c_7;
+constexpr auto cb_max_value_after_reduction = tt::CBIndex::c_8;
+constexpr auto cb_exp_sum_before_reduction = tt::CBIndex::c_9;
+constexpr auto cb_exp_sum_after_reduction = tt::CBIndex::c_10;
+constexpr auto cb_output = tt::CBIndex::c_11;
+
+constexpr uint32_t onetile = 1;
+
+#ifdef DO_MASK_W
+constexpr bool do_mask_w = true;
+#else
+constexpr bool do_mask_w = false;
+#endif
+
+#ifdef EVERYTHING_FITS_IN_L1
+
+void find_max_value_in_row() {
+    const uint32_t max_value_register = 0;
+    const uint32_t tile_register = 1U;
+    cb_wait_front(cb_input, Wt);  // wait until reader kernel has written Wt tiles to input buffer
+    cb_reserve_back(cb_max_value_before_reduction, onetile);
+    tile_regs_acquire();
+    for (uint32_t col = 0; col < Wt; ++col) {
+        auto working_register = col == 0 ? max_value_register : tile_register;
+        copy_tile_init(cb_input);
+        copy_tile(cb_input, /* tile_idx */ col, /* register_idx */ working_register);
+
+        if constexpr (do_mask_w) {
+            if (col + 1 == Wt) {
+                // this is limitation of the function mask_tile
+                // mask tile currently does not work for mask register that is not next to data register
+                const uint32_t mask_register = working_register + 1U;  // mask register should be next to data register
+
+                // the next 4 lines are important because we overwrite what's in the trash padding.
+                // it's possible that the padding contains a NaN, and operations like NaN + (-inf) = NaN,
+                // instead of the expected -inf. similarly, -inf * 0 = NaN.
+                copy_tile_init(cb_mask);
+                copy_tile(cb_mask, /* tile_idx */ 0, /* register idx */ mask_register);
+                mask_tile_init();
+                mask_tile(working_register, mask_register);  // mask should be next to tile register.
+
+                copy_tile_init(cb_max_mask);
+                copy_tile(cb_max_mask, /* tile_idx */ 0, /* register idx */ mask_register);
+
+                add_binary_tile_init();
+                add_binary_tile(working_register, mask_register);
+            }
+        }
+
+        if (col > 0) {
+            max_tile_init();
+            max_tile(max_value_register, tile_register);
+        }
+    }
+    tile_regs_commit();
+
+    tile_regs_wait();
+    pack_reconfig_data_format(cb_max_value_before_reduction);
+    pack_tile(max_value_register, cb_max_value_before_reduction);
+    tile_regs_release();
+    cb_push_back(cb_max_value_before_reduction, onetile);
+}
+
+#else
+
+void find_max_value_in_row() {
+    cb_reserve_back(cb_max_value_before_reduction, onetile);
+    tile_regs_acquire();
+    const uint32_t max_value_register = 0;
+    const uint32_t tile_register = 1U;
+    for (uint32_t col = 0; col < Wt;) {
+        cb_wait_front(cb_input, block_size);  // wait until reader kernel has written block_size tiles to input buffer
+        for (uint32_t block_idx = 0; block_idx < block_size; ++block_idx, ++col) {
+            auto working_register = col == 0 ? max_value_register : tile_register;
+            copy_tile_init(cb_input);
+            copy_tile(cb_input, /* tile_idx */ block_idx, /* register_idx */ working_register);
+
+            if constexpr (do_mask_w) {
+                if (col + 1 == Wt) {
+                    // this is limitation of the function mask_tile
+                    // mask tile currently does not work for mask register that is not next to data register
+                    const uint32_t mask_register =
+                        working_register + 1U;  // mask register should be next to data register
+
+                    // the next 4 lines are important because we overwrite what's in the trash padding.
+                    // it's possible that the padding contains a NaN, and operations like NaN + (-inf) = NaN,
+                    // instead of the expected -inf. similarly, -inf * 0 = NaN.
+                    copy_tile_init(cb_mask);
+                    copy_tile(cb_mask, /* tile_idx */ 0, /* register idx */ mask_register);
+                    mask_tile_init();
+                    mask_tile(working_register, mask_register);  // mask should be next to tile register.
+
+                    copy_tile_init(cb_max_mask);
+                    copy_tile(cb_max_mask, /* tile_idx */ 0, /* register idx */ mask_register);
+
+                    add_binary_tile_init();
+                    add_binary_tile(working_register, mask_register);
+                }
+            }
+
+            if (col > 0) {
+                max_tile_init();
+                max_tile(max_value_register, tile_register);
+            }
+        }
+        cb_pop_front(cb_input, block_size);  // delete block_size tiles from input buffer
+    }
+    tile_regs_commit();
+
+    tile_regs_wait();
+    pack_reconfig_data_format(cb_max_value_before_reduction);
+    pack_tile(max_value_register, cb_max_value_before_reduction);
+    tile_regs_release();
+    cb_push_back(cb_max_value_before_reduction, onetile);
+}
+
+#endif
+
+void reduce_max_value() {
+    // reduce max value and push it to cb_max_value_after_reduction
+    cb_wait_front(cb_max_value_before_reduction, onetile);
+    cb_reserve_back(cb_max_value_after_reduction, onetile);
+
+    const uint32_t reduction_register = 0;
+    tile_regs_acquire();
+    reconfig_data_format(cb_max_value_before_reduction, cb_scaler);
+    reduce_init_delta<false, PoolType::MAX, ReduceDim::REDUCE_ROW>(
+        cb_max_value_before_reduction, cb_scaler, cb_max_value_after_reduction);
+    reduce_tile<PoolType::MAX, ReduceDim::REDUCE_ROW>(
+        cb_max_value_before_reduction,
+        cb_scaler,
+        /* tile_idx */ 0,
+        /* tile_idx */ 0,
+        reduction_register);
+    reduce_revert_delta<ReduceDim::REDUCE_ROW>(cb_max_value_before_reduction);
+    tile_regs_commit();
+
+    tile_regs_wait();
+
+    pack_reconfig_data_format(cb_max_value_after_reduction);
+    pack_tile(reduction_register, cb_max_value_after_reduction);
+    tile_regs_release();
+
+    cb_push_back(cb_max_value_after_reduction, onetile);   // push tile after reduction
+    cb_pop_front(cb_max_value_before_reduction, onetile);  // pop tile before reduction
+}
+
+#ifdef EVERYTHING_FITS_IN_L1
+
+// calculate sum of exp(x - max(x))
+void calculate_sum_exp_x() {
+    cb_wait_front(cb_max_value_after_reduction, onetile);  // wait until we get max value in each row
+
+    // run through all tiles in row
+    const uint32_t accum_register = 0;
+    const uint32_t tile_register = 1U;
+
+    cb_reserve_back(cb_exp_sum_before_reduction, onetile);
+    tile_regs_acquire();
+
+    const uint32_t max_value_register = 3U;
+    unary_bcast_init<BroadcastType::COL>(cb_max_value_after_reduction, cb_max_value_after_reduction);
+    unary_bcast<BroadcastType::COL>(
+        cb_max_value_after_reduction, /* tile idx */ 0, /* reg tile idx */ max_value_register);
+    for (uint32_t col = 0; col < Wt; ++col) {
+        auto working_register = col == 0 ? accum_register : tile_register;
+
+        copy_tile_init(cb_input);
+        copy_tile(cb_input, /* tile_idx */ col, /* register_idx */ working_register);
+
+        sub_binary_tile_init();
+        sub_binary_tile(working_register, max_value_register);  // subtract max value from each tile
+
+        exp_tile_init();
+        exp_tile</* approx */ false>(working_register);  // calculate exp for each tile in tile register
+
+        if constexpr (do_mask_w) {
+            if (col + 1 == Wt) {
+                // this is limitation of the function mask_tile
+                // mask tile currently does not work for mask register that is not next to data register
+                const uint32_t mask_register = working_register + 1U;  // mask register should be next to data register
+                copy_tile_init(cb_mask);
+                copy_tile(cb_mask, /* tile_idx */ 0, /* register idx */ mask_register);
+
+                mask_tile_init();
+                mask_tile(working_register, mask_register);  // mask should be next to tile register
+            }
+        }
+
+        if (col > 0) {
+            add_binary_tile_init();
+            add_binary_tile(accum_register, working_register);
+        }
+    }
+    tile_regs_commit();
+
+    // packing to cb_exp_sum_before_reduction
+    tile_regs_wait();
+    pack_reconfig_data_format(cb_exp_sum_before_reduction);
+    pack_tile(accum_register, cb_exp_sum_before_reduction);
+    tile_regs_release();
+    cb_push_back(cb_exp_sum_before_reduction, onetile);
+}
+
+#else
+
+// calculate sum of exp(x - max(x))
+void calculate_sum_exp_x() {
+    // wait until we get max value in each row
+    cb_wait_front(cb_max_value_after_reduction, onetile);  // wait until we get max value in each row
+    cb_reserve_back(cb_exp_sum_before_reduction, onetile);
+    // run through all tiles in row
+    const uint32_t accum_register = 0;
+    const uint32_t tile_register = 1U;
+
+    tile_regs_acquire();
+
+    const uint32_t max_value_register = 3U;
+    unary_bcast_init<BroadcastType::COL>(cb_max_value_after_reduction, cb_max_value_after_reduction);
+    unary_bcast<BroadcastType::COL>(
+        cb_max_value_after_reduction, /* tile idx */ 0, /* reg tile idx */ max_value_register);
+    for (uint32_t col = 0; col < Wt;) {
+        cb_wait_front(cb_input, block_size);  // wait until reader kernel has written block_size tiles to input buffer
+        reconfig_data_format(cb_max_value_after_reduction, cb_input);
+        for (uint32_t block_idx = 0; block_idx < block_size; ++block_idx, ++col) {
+            auto working_register = col == 0 ? accum_register : tile_register;
+
+            copy_tile_init(cb_input);
+            copy_tile(cb_input, /* tile_idx */ block_idx, /* register_idx */ working_register);
+
+            sub_binary_tile_init();
+            sub_binary_tile(working_register, max_value_register);  // subtract max value from each tile
+
+            exp_tile_init();
+            exp_tile</* approx */ false>(working_register);  // calculate exp for each tile in tile register
+
+            if constexpr (do_mask_w) {
+                if (col + 1 == Wt) {
+                    // this is limitation of the function mask_tile
+                    // mask tile currently does not work for mask register that is not next to data register
+                    const uint32_t mask_register =
+                        working_register + 1U;  // mask register should be next to data register
+                    copy_tile_init(cb_mask);
+                    copy_tile(cb_mask, /* tile_idx */ 0, /* register idx */ mask_register);
+
+                    mask_tile_init();
+                    mask_tile(working_register, mask_register);  // mask should be next to tile register
+                }
+            }
+
+            if (col > 0) {
+                add_binary_tile_init();
+                add_binary_tile(accum_register, working_register);
+            }
+        }
+        cb_pop_front(cb_input, block_size);
+    }
+    tile_regs_commit();
+
+    // packing to cb_exp_sum_before_reduction
+    tile_regs_wait();
+    pack_reconfig_data_format(cb_exp_sum_before_reduction);
+    pack_tile(accum_register, cb_exp_sum_before_reduction);
+    tile_regs_release();
+    cb_push_back(cb_exp_sum_before_reduction, onetile);
+}
+
+#endif
+
+void reduce_log_sum_exp_x() {
+    // reduce exp sum to cb_exp_sum_after_reduction
+    cb_wait_front(cb_exp_sum_before_reduction, onetile);
+    cb_reserve_back(cb_exp_sum_after_reduction, onetile);
+
+    tile_regs_acquire();
+    const uint32_t reduction_register = 0;
+    reconfig_data_format(cb_exp_sum_before_reduction, cb_scaler);
+    reduce_init_delta<false, PoolType::SUM, ReduceDim::REDUCE_ROW>(
+        cb_exp_sum_before_reduction, cb_scaler, cb_exp_sum_after_reduction);
+    reduce_tile<PoolType::SUM, ReduceDim::REDUCE_ROW>(
+        cb_exp_sum_before_reduction,
+        cb_scaler,
+        /* tile_idx */ 0,
+        /* tile_idx */ 0,
+        /* reduction_register */ reduction_register);
+    reduce_revert_delta<ReduceDim::REDUCE_ROW>(cb_exp_sum_before_reduction);
+
+    // log(sum(exp(x - max(x))))
+    log_tile_init();
+    log_tile(reduction_register);
+    tile_regs_commit();
+
+    tile_regs_wait();
+    pack_reconfig_data_format(cb_exp_sum_after_reduction);
+    pack_tile(/* tile_idx */ reduction_register, cb_exp_sum_after_reduction);
+    tile_regs_release();
+    cb_push_back(cb_exp_sum_after_reduction, onetile);
+    cb_pop_front(cb_exp_sum_before_reduction, onetile);
+}
+
+void MAIN {
+    if constexpr (do_mask_w) {
+        cb_wait_front(cb_mask, onetile);
+        cb_wait_front(cb_max_mask, onetile);
+    }
+    cb_wait_front(cb_scaler, onetile);
+
+    init_sfpu(cb_input, cb_output);
+    binary_op_init_common(cb_input, cb_target_logits, cb_output);
+
+    for (uint32_t row = 0; row < num_rows_per_core; ++row) {
+        find_max_value_in_row();  // find max value in each row
+        reduce_max_value();       // reduce max value to cb_max_value_after_reduction
+
+        calculate_sum_exp_x();   // calculate sum of exp(x - max(x))
+        reduce_log_sum_exp_x();  // reduce sum(exp(x - max(x))), take log and push to cb_exp_sum_after_reduction
+
+        cb_wait_front(cb_exp_sum_after_reduction, onetile);  //  wait log(sum(exp(x - max(x)))
+        cb_wait_front(cb_target_logits, onetile);
+        cb_reserve_back(cb_output, onetile);  // reserve onetile tiles in output buffer
+
+        const uint32_t result_register = 0;
+        tile_regs_acquire();
+        // calculate result as: -target_logits + max_value + log(sum(exp(x - max(x))))
+        // -(target_logits - max_value) == -target_logits + max_value
+        copy_tile_init(cb_target_logits);
+        copy_tile(cb_target_logits, /* tile_idx */ 0, /* register_idx */ result_register);
+
+        negative_tile_init();
+        negative_tile(result_register);
+
+        reconfig_data_format(cb_max_value_after_reduction, cb_exp_sum_after_reduction);
+        add_tiles_init(cb_max_value_after_reduction, cb_exp_sum_after_reduction, /* acc_to_dest */ true);
+        add_tiles(
+            cb_max_value_after_reduction,
+            cb_exp_sum_after_reduction,
+            /* tile_idx */ 0,
+            /* tile_idx */ 0,
+            result_register);
+
+        tile_regs_commit();
+
+        tile_regs_wait();
+        pack_reconfig_data_format(cb_output);
+        pack_tile(result_register, cb_output);
+        tile_regs_release();
+        cb_push_back(cb_output, onetile);
+
+        cb_pop_front(cb_max_value_after_reduction, onetile);  // pop tile after reduction
+        cb_pop_front(cb_exp_sum_after_reduction, onetile);    // pop tile after reduction
+        cb_pop_front(cb_target_logits, onetile);              // pop logits tile
+
+#ifdef EVERYTHING_FITS_IN_L1
+        cb_pop_front(cb_input, Wt);  // pop Wt tiles from input buffer
+#endif
+    }
+
+    // pop scaler and masks
+    if constexpr (do_mask_w) {
+        cb_pop_front(cb_mask, onetile);
+        cb_pop_front(cb_max_mask, onetile);
+    }
+    cb_pop_front(cb_scaler, onetile);
+}
+
+}  // namespace NAMESPACE
diff --git a/tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/kernels/dataflow/reader_cross_entropy_fw_interleaved_start_id.cpp b/tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/kernels/dataflow/reader_cross_entropy_fw_interleaved_start_id.cpp
new file mode 100644
index 000000000000..359e62312ada
--- /dev/null
+++ b/tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/kernels/dataflow/reader_cross_entropy_fw_interleaved_start_id.cpp
@@ -0,0 +1,215 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include <dataflow_api_addrgen.h>
+
+#include <cstdint>
+#include <cstring>
+
+#include "dataflow_api.h"
+#include "debug/dprint.h"
+#include "debug/dprint_pages.h"
+
+constexpr uint32_t FACE_HEIGHT = 16;
+constexpr uint32_t FACE_WIDTH = 16;
+constexpr uint32_t TILE_HEIGHT = 32;
+constexpr uint32_t TILE_WIDTH = 32;
+constexpr uint32_t TARGET_INDEXES_ELEMENTS_NUMBER = 32;
+
+// calculate page and offset for target indexes
+std::pair<uint32_t, uint32_t> get_page_and_offset(uint32_t tiled_row, uint32_t tiled_H) {
+    uint32_t n = tiled_row / tiled_H;
+    uint32_t h = (tiled_row % tiled_H) * 32;
+
+    uint32_t page = n;
+    uint32_t offset = h * sizeof(uint32_t);
+    return {page, offset};
+}
+
+uint32_t get_tilized_idx(uint32_t h, uint32_t w) {
+    // Get local coordinates within the tile
+    uint32_t local_row = h % TILE_HEIGHT;
+    uint32_t local_col = w % TILE_WIDTH;
+
+    // Determine the index offset based on which quadrant we're in
+    uint32_t offset = 0;
+
+    // If we're in the right half (columns beyond FACE_WIDTH)
+    if (local_col >= FACE_WIDTH) {
+        local_col -= FACE_WIDTH;
+        offset += FACE_HEIGHT * FACE_WIDTH;  // Right face offset
+    }
+
+    // If we're in the bottom half (rows beyond FACE_WIDTH)
+    if (local_row >= FACE_WIDTH) {
+        local_row -= FACE_WIDTH;
+        offset += FACE_HEIGHT * TILE_WIDTH;  // Bottom face offset
+    }
+
+    // Final index within the tile
+    uint32_t index = offset + local_row * FACE_WIDTH + local_col;
+    return index;
+}
+
+void kernel_main() {
+    uint32_t runtime_args_counter = 0U;
+    uint32_t input_address = get_arg_val<uint32_t>(runtime_args_counter++);        // input buffer address
+    uint32_t target_address = get_arg_val<uint32_t>(runtime_args_counter++);       // target buffer address
+    uint32_t num_rows_to_process = get_arg_val<uint32_t>(runtime_args_counter++);  // rows to process in this kernel
+    uint32_t start_row =
+        get_arg_val<uint32_t>(runtime_args_counter++);  // pre calculated num_rows_written in program factory
+
+    constexpr uint32_t cb_input_idx = tt::CBIndex::c_0;
+    constexpr uint32_t cb_target_idx = tt::CBIndex::c_1;
+    constexpr uint32_t cb_mask_idx = tt::CBIndex::c_2;
+    constexpr uint32_t cb_max_mask_idx = tt::CBIndex::c_3;
+    constexpr uint32_t cb_scaler_idx = tt::CBIndex::c_4;  // used for reduction
+    constexpr uint32_t cb_input_tile = tt::CBIndex::c_5;
+    constexpr uint32_t cb_target_logits = tt::CBIndex::c_6;
+
+    constexpr uint32_t block_size = get_compile_time_arg_val(0);
+    constexpr uint32_t Wt = get_compile_time_arg_val(1);
+    constexpr uint32_t mask_w = get_compile_time_arg_val(2);
+    constexpr uint32_t target_indexes_page_size = get_compile_time_arg_val(3);
+    constexpr uint32_t tiled_H = get_compile_time_arg_val(4);
+    constexpr uint32_t target_indexes_read_page_size = get_compile_time_arg_val(5);
+
+    constexpr uint32_t onetile = 1U;
+#ifdef DO_MASK_W
+    constexpr bool do_mask_w = true;
+#else
+    constexpr bool do_mask_w = false;
+#endif
+
+    // generate scaler and mask tile
+    cb_reserve_back(cb_scaler_idx, onetile);
+    uint16_t* scaler_ptr = reinterpret_cast<uint16_t*>(get_write_ptr(cb_scaler_idx));  // write scalar tile
+
+    uint16_t* mask_ptr = nullptr;
+    uint16_t* max_mask_ptr = nullptr;
+    if constexpr (do_mask_w) {
+        cb_reserve_back(cb_mask_idx, onetile);
+        cb_reserve_back(cb_max_mask_idx, onetile);
+        mask_ptr = reinterpret_cast<uint16_t*>(get_write_ptr(cb_mask_idx));          // write mask tile
+        max_mask_ptr = reinterpret_cast<uint16_t*>(get_write_ptr(cb_max_mask_idx));  // write max mask tile
+    }
+
+    constexpr uint16_t one = 0x00003F80;  // (bfloat16)1.0 -> uint16_t
+    constexpr uint16_t zero = 0x0;
+    constexpr uint16_t minus_inf = 0xFF80;  // (bfloat16)-inf -> uint16_t
+    for (uint32_t face = 0; face < 4; ++face) {
+        uint32_t offset = (face & 1U) << 4U;
+        for (uint32_t h = 0; h < 16; ++h) {
+            for (uint32_t w = 0; w < 16; ++w) {
+                if constexpr (do_mask_w) {
+                    *mask_ptr++ = (offset + w < mask_w) ? one : zero;  // how to create the proper mask?
+                    *max_mask_ptr++ = (offset + w < mask_w) ? zero : minus_inf;
+                }
+
+                *scaler_ptr++ = one;
+            }
+        }
+    }
+    if constexpr (do_mask_w) {
+        cb_push_back(cb_mask_idx, onetile);
+        cb_push_back(cb_max_mask_idx, onetile);
+    }
+    cb_push_back(cb_scaler_idx, onetile);
+
+    const uint32_t tile_bytes = get_tile_size(cb_input_idx);
+    const DataFormat data_format = get_dataformat(cb_input_idx);
+
+    const InterleavedAddrGenFast</* is_dram */ true> input_address_generator = {
+        .bank_base_address = input_address, .page_size = tile_bytes, .data_format = data_format};
+
+    const InterleavedAddrGen</* is_dram */ true> target_indexes_address_generator = {
+        .bank_base_address = target_address, .page_size = target_indexes_page_size};
+
+    for (uint32_t i = 0; i < num_rows_to_process; ++i) {
+        // calculate the address of the first tile in the row
+        // start_row is the number of rows already processed in other cores
+        uint32_t idx = (start_row + i) * Wt;  // (take already processed rows + current row)*Wt(number of tiles in row)
+
+        // read target indexes
+        cb_reserve_back(cb_target_idx, onetile);
+        uint32_t l1_target_indexes_write_addr =
+            get_write_ptr(cb_target_idx);  // get the address of the first tile in the target buffer
+
+        auto [page, offset] = get_page_and_offset(start_row + i, tiled_H);
+
+        auto noc_async_target_indexes_page_addr = get_noc_addr(page, target_indexes_address_generator, offset);
+        noc_async_read(
+            noc_async_target_indexes_page_addr,
+            l1_target_indexes_write_addr,
+            target_indexes_read_page_size);  // read the page from the target buffer
+        noc_async_read_barrier();            // wait until all tiles are read
+
+        cb_reserve_back(cb_input_tile, onetile);
+        cb_reserve_back(cb_target_logits, onetile);
+
+        uint32_t l1_input_tile_write_addr = get_write_ptr(cb_input_tile);
+        uint32_t l1_target_logits_write_addr = get_write_ptr(cb_target_logits);
+
+        auto target_indexes_l1_ptr = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_read_ptr(cb_target_idx));
+        for (uint32_t h = 0; h < TILE_HEIGHT; ++h) {
+            uint32_t target_value_idx = h;  // only first row of the tile is used
+
+            uint32_t target_value = target_indexes_l1_ptr[target_value_idx];
+
+            noc_async_read_tile(idx + (target_value / TILE_WIDTH), input_address_generator, l1_input_tile_write_addr);
+            noc_async_read_barrier();
+
+            auto read_input_tile_l1_ptr = reinterpret_cast<volatile tt_l1_ptr uint16_t*>(get_read_ptr(cb_input_tile));
+
+            auto target_logits_write_l1_ptr =
+                reinterpret_cast<volatile tt_l1_ptr uint16_t*>(l1_target_logits_write_addr);
+
+            auto idx_inside_tile = get_tilized_idx(h, target_value);  // only first row of the tile is used
+            uint32_t inplace_idx = get_tilized_idx(h, 0);
+            target_logits_write_l1_ptr[inplace_idx] = read_input_tile_l1_ptr[idx_inside_tile];
+        }
+
+        cb_push_back(cb_target_logits, onetile);
+
+#ifdef EVERYTHING_FITS_IN_L1
+        // read input buffer
+        cb_reserve_back(cb_input_idx, Wt);  // reserve Wt tiles in input buffer ==  wait until cb will has Wt tiles
+        uint32_t l1_write_addr = get_write_ptr(cb_input_idx);  // get the address of the first tile in the input buffer
+
+        for (uint32_t j = 0; j < Wt; ++j) {
+            noc_async_read_tile(
+                idx + j, input_address_generator, l1_write_addr);  // read the tile from the input buffer
+            l1_write_addr += tile_bytes;                           // move to the next tile
+        }
+        noc_async_read_barrier();        // wait until all tiles are read
+        cb_push_back(cb_input_idx, Wt);  // push the tile to the back of the input buffer
+
+#else
+        // read input buffer by blocks to calculate max value in row
+        for (uint32_t j = 0; j < Wt; j += block_size) {
+            cb_reserve_back(cb_input_idx, block_size);
+            uint32_t l1_write_addr = get_write_ptr(cb_input_idx);
+            for (uint32_t block_idx = 0; block_idx < block_size; ++block_idx) {
+                noc_async_read_tile(idx + j + block_idx, input_address_generator, l1_write_addr);
+                l1_write_addr += tile_bytes;
+            }
+            noc_async_read_barrier();
+            cb_push_back(cb_input_idx, block_size);
+        }
+
+        // read input buffer by blocks to calculate sum(exp(x - max(x))) in row
+        for (uint32_t j = 0; j < Wt; j += block_size) {
+            cb_reserve_back(cb_input_idx, block_size);
+            uint32_t l1_write_addr = get_write_ptr(cb_input_idx);
+            for (uint32_t block_idx = 0; block_idx < block_size; ++block_idx) {
+                noc_async_read_tile(idx + j + block_idx, input_address_generator, l1_write_addr);
+                l1_write_addr += tile_bytes;
+            }
+            noc_async_read_barrier();
+            cb_push_back(cb_input_idx, block_size);
+        }
+
+#endif
+    }
+}
diff --git a/tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/kernels/dataflow/writer_cross_entropy_fw_interleaved_start_id.cpp b/tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/kernels/dataflow/writer_cross_entropy_fw_interleaved_start_id.cpp
new file mode 100644
index 000000000000..06a1349d0e11
--- /dev/null
+++ b/tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/kernels/dataflow/writer_cross_entropy_fw_interleaved_start_id.cpp
@@ -0,0 +1,36 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include "dataflow_api.h"
+
+void kernel_main() {
+    uint32_t runtime_args_counter = 0;
+    uint32_t output_addr = get_arg_val<uint32_t>(runtime_args_counter++);
+    uint32_t num_rows_to_process = get_arg_val<uint32_t>(runtime_args_counter++);
+    uint32_t start_row = get_arg_val<uint32_t>(runtime_args_counter++);
+
+    constexpr uint32_t cb_output_idx = tt::CBIndex::c_11;
+
+    constexpr uint32_t block_size = get_compile_time_arg_val(0);
+    constexpr uint32_t Wt = get_compile_time_arg_val(1);  // number of tiles in inner dimension
+
+    constexpr uint32_t onetile = 1U;
+
+    const uint32_t tile_bytes = get_tile_size(cb_output_idx);
+    const DataFormat data_format = get_dataformat(cb_output_idx);
+
+    const InterleavedAddrGenFast</* is dram */ true> output_addr_generator = {
+        .bank_base_address = output_addr, .page_size = tile_bytes, .data_format = data_format};
+
+    uint32_t end_row = start_row + num_rows_to_process;
+
+    for (uint32_t r = start_row; r < end_row; r++) {
+        cb_wait_front(cb_output_idx, onetile);
+        uint32_t l1_read_addr = get_read_ptr(cb_output_idx);
+
+        noc_async_write_tile(r, output_addr_generator, l1_read_addr);
+        noc_async_write_barrier();
+        cb_pop_front(cb_output_idx, onetile);
+    }
+}
diff --git a/tt-train/sources/ttml/metal/ops/rmsnorm_fw/rmsnorm_fw.cpp b/tt-train/sources/ttml/metal/ops/rmsnorm_fw/rmsnorm_fw.cpp
index 472dd47929fb..77b3fae3c227 100644
--- a/tt-train/sources/ttml/metal/ops/rmsnorm_fw/rmsnorm_fw.cpp
+++ b/tt-train/sources/ttml/metal/ops/rmsnorm_fw/rmsnorm_fw.cpp
@@ -1,4 +1,3 @@
-
 // SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
 //
 // SPDX-License-Identifier: Apache-2.0
diff --git a/tt-train/sources/ttml/metal/ops/rmsnorm_fw/rmsnorm_fw.hpp b/tt-train/sources/ttml/metal/ops/rmsnorm_fw/rmsnorm_fw.hpp
index f8f16084422f..1751908c4886 100644
--- a/tt-train/sources/ttml/metal/ops/rmsnorm_fw/rmsnorm_fw.hpp
+++ b/tt-train/sources/ttml/metal/ops/rmsnorm_fw/rmsnorm_fw.hpp
@@ -1,4 +1,3 @@
-
 // SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
 //
 // SPDX-License-Identifier: Apache-2.0
diff --git a/tt-train/tests/ops/cross_entropy_fw_op_test.cpp b/tt-train/tests/ops/cross_entropy_fw_op_test.cpp
new file mode 100644
index 000000000000..6deb69297a29
--- /dev/null
+++ b/tt-train/tests/ops/cross_entropy_fw_op_test.cpp
@@ -0,0 +1,287 @@
+
+// SPDX-FileCopyrightText: (c) 2025 Tenstorrent AI ULC
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include <gtest/gtest.h>
+#include <sys/types.h>
+
+#include <cassert>
+#include <core/ttnn_all_includes.hpp>
+#include <cstddef>
+#include <cstdint>
+#include <ttnn/operations/reduction/generic/generic_reductions.hpp>
+
+#include "autograd/auto_context.hpp"
+#include "core/tt_tensor_utils.hpp"
+#include "metal/operations.hpp"
+
+class CrossEntropyForwardTest : public ::testing::Test {
+protected:
+    void SetUp() override {
+        ttml::autograd::ctx().open_device();
+    }
+
+    void TearDown() override {
+        ttml::autograd::ctx().close_device();
+    }
+};
+
+xt::xarray<float> calculate_cross_entropy_loss(const xt::xarray<float>& input, const xt::xarray<uint32_t>& target) {
+    const uint32_t N = target.shape(0);
+    const uint32_t C = 1U;
+    const uint32_t H = target.shape(1);
+    const uint32_t W = 1U;
+    xt::xarray<float> target_inputs = xt::zeros<float>({N, C, H, W});
+
+    for (size_t n = 0; n < N; ++n) {
+        for (size_t h = 0; h < H; ++h) {
+            size_t class_index = target(n, h);
+            target_inputs(n, 0, h, 0) = input(n, 0, h, class_index);
+        }
+    }
+
+    xt::xarray<float> max_input = xt::amax(input, -1, xt::keep_dims);
+    xt::xarray<float> shifted_input = input - max_input;
+    xt::xarray<float> log_exp_sum_test = xt::log(xt::sum(xt::exp(shifted_input), -1, xt::keep_dims));
+    xt::xarray<float> result = -target_inputs + max_input + log_exp_sum_test;
+    return result;
+}
+
+TEST_F(CrossEntropyForwardTest, CrossEntropyForward_Small_Forward) {
+    using namespace ttml;
+
+    const uint32_t N = 1, C = 1, H = 1, W = 8;
+
+    xt::xarray<float> input_tensor = {{{{1.F, 2.F, 3.F, 4.F, 1.F, 2.F, 3.F, 4.F}}}};
+    auto input = core::from_xtensor(input_tensor, &autograd::ctx().get_device());
+    std::cout << "Input Logits:\n";
+    input.print();
+
+    xt::xarray<uint32_t> target_tensor = xt::zeros<uint32_t>({N, H});
+    target_tensor(0, 0) = 1U;
+    auto target = core::from_xtensor<uint32_t, ttnn::DataType::UINT32>(
+        target_tensor, &autograd::ctx().get_device(), ttnn::Layout::ROW_MAJOR);
+    std::cout << "Input Target Indexes:\n";
+    target.print();
+
+    auto result = ttml::metal::cross_entropy_fw(input, target);
+    std::cout << "CrossEntropyForward_Test:\nResult:\n";
+    result.print();
+
+    auto expected_result = calculate_cross_entropy_loss(input_tensor, target_tensor);
+    auto expected_result_print = core::from_xtensor(expected_result, &autograd::ctx().get_device());
+    std::cout << "Expected Result:\n";
+    expected_result_print.print();
+
+    // Check if the result is close to the expected result
+    auto result_xtensor = core::to_xtensor(result);
+    assert((result_xtensor.shape() == expected_result.shape()));
+    EXPECT_TRUE(xt::allclose(result_xtensor, expected_result, 3e-2F, 1e-2F));
+}
+
+TEST_F(CrossEntropyForwardTest, CrossEntropyForward_Negetive_Values) {
+    using namespace ttml;
+
+    const uint32_t N = 1, C = 1, H = 2, W = 4;
+
+    xt::xarray<float> input_tensor = {{{{-100.F, -101.F, -102.F, -103.F}, {-5.01F, -5.02F, -0.3F, -7.F}}}};
+    auto input = core::from_xtensor(input_tensor, &autograd::ctx().get_device());
+    std::cout << "Input Logits:\n";
+    input.print();
+
+    xt::xarray<uint32_t> target_tensor = xt::zeros<uint32_t>({N, H});
+    target_tensor(0, 0) = 0;
+    target_tensor(0, 1) = 2U;
+    auto target = core::from_xtensor<uint32_t, ttnn::DataType::UINT32>(
+        target_tensor, &autograd::ctx().get_device(), ttnn::Layout::ROW_MAJOR);
+    std::cout << "Input Target Indexes:\n";
+    target.print();
+
+    auto result = ttml::metal::cross_entropy_fw(input, target);
+    std::cout << "CrossEntropyForward_Test:\nResult:\n";
+    result.print();
+
+    auto expected_result = calculate_cross_entropy_loss(input_tensor, target_tensor);
+    auto expected_result_print = core::from_xtensor(expected_result, &autograd::ctx().get_device());
+    std::cout << "Expected Result:\n";
+    expected_result_print.print();
+
+    // Check if the result is close to the expected result
+    auto result_xtensor = core::to_xtensor(result);
+    assert((result_xtensor.shape() == expected_result.shape()));
+    EXPECT_TRUE(xt::allclose(result_xtensor, expected_result, 3e-2F, 1e-2F));
+}
+
+TEST_F(CrossEntropyForwardTest, CrossEntropyForward_Batch) {
+    using namespace ttml;
+
+    const uint32_t N = 2U, C = 1U, H = 91U, W = 157U;
+    const auto shape = ttnn::SmallVector<uint32_t>{N, C, H, W};
+
+    std::random_device rd;
+    std::mt19937 gen(42);
+    xt::xarray<float> input_tensor = xt::random::rand<float>({N, C, H, W}, -10.0F, 10.0F, gen);
+    xt::xarray<uint32_t> target_tensor = xt::zeros<uint32_t>({N, H});
+
+    std::uniform_int_distribution<uint32_t> class_dist(0, W - 1);
+    for (uint32_t n = 0; n < N; ++n) {
+        for (uint32_t h = 0; h < H; ++h) {
+            uint32_t true_class = class_dist(gen);
+            target_tensor(n, h) = true_class;
+        }
+    }
+
+    auto input = core::from_xtensor(input_tensor, &autograd::ctx().get_device());
+    std::cout << "Input Logits:\n";
+    input.print();
+
+    auto target = core::from_xtensor<uint32_t, ttnn::DataType::UINT32>(
+        target_tensor, &autograd::ctx().get_device(), ttnn::Layout::ROW_MAJOR);
+    std::cout << "Input Target Indexes:\n";
+    target.print();
+
+    auto result = ttml::metal::cross_entropy_fw(input, target);
+    std::cout << "CrossEntropyForward_Test:\nResult:\n";
+    result.print();
+
+    auto expected_result = calculate_cross_entropy_loss(input_tensor, target_tensor);
+
+    auto expected_result_print = core::from_xtensor(expected_result, &autograd::ctx().get_device());
+    std::cout << "Expected Result:\n";
+    expected_result_print.print();
+
+    // Check if the result is close to the expected result
+    auto result_xtensor = core::to_xtensor(result);
+    assert((result_xtensor.shape() == expected_result.shape()));
+    EXPECT_TRUE(xt::allclose(result_xtensor, expected_result, 3e-2F, 1e-2F));
+}
+
+TEST_F(CrossEntropyForwardTest, CrossEntropyForward_Large_Batch) {
+    using namespace ttml;
+
+    const uint32_t N = 64U, C = 1U, H = 1017U, W = 1018U;
+    const auto shape = ttnn::SmallVector<uint32_t>{N, C, H, W};
+
+    std::random_device rd;
+    std::mt19937 gen(42);
+    xt::xarray<float> input_tensor = xt::random::rand<float>({N, C, H, W}, -10.0F, 10.0F, gen);
+    xt::xarray<uint32_t> target_tensor = xt::zeros<uint32_t>({N, H});
+
+    std::uniform_int_distribution<uint32_t> class_dist(0, W - 1);
+    for (uint32_t n = 0; n < N; ++n) {
+        for (uint32_t h = 0; h < H; ++h) {
+            uint32_t true_class = class_dist(gen);
+            target_tensor(n, h) = true_class;
+        }
+    }
+
+    auto input = core::from_xtensor(input_tensor, &autograd::ctx().get_device());
+    std::cout << "Input Logits:\n";
+    input.print();
+
+    auto target = core::from_xtensor<uint32_t, ttnn::DataType::UINT32>(
+        target_tensor, &autograd::ctx().get_device(), ttnn::Layout::ROW_MAJOR);
+    std::cout << "Input Target Indexes:\n";
+    target.print();
+
+    auto result = ttml::metal::cross_entropy_fw(input, target);
+    std::cout << "CrossEntropyForward_Test:\nResult:\n";
+    result.print();
+
+    auto expected_result = calculate_cross_entropy_loss(input_tensor, target_tensor);
+    auto expected_result_print = core::from_xtensor(expected_result, &autograd::ctx().get_device());
+    std::cout << "Expected Result:\n";
+    expected_result_print.print();
+
+    // Check if the result is close to the expected result
+    auto result_xtensor = core::to_xtensor(result);
+    assert((result_xtensor.shape() == expected_result.shape()));
+    EXPECT_TRUE(xt::allclose(result_xtensor, expected_result, 3e-2F, 1e-2F));
+}
+
+TEST_F(CrossEntropyForwardTest, CrossEntropyForward_Large_Forward) {
+    using namespace ttml;
+
+    const uint32_t N = 1U, C = 1U, H = 1U, W = 65536U;
+    const auto shape = ttnn::SmallVector<size_t>{N, C, H, W};
+
+    std::random_device rd;
+    std::mt19937 gen(42);
+    xt::xarray<float> input_tensor = xt::random::rand<float>({N, C, H, W}, -10.0F, 10.0F, gen);
+    xt::xarray<uint32_t> target_tensor = xt::zeros<uint32_t>({N, H});
+
+    std::uniform_int_distribution<uint32_t> class_dist(0, W - 1);
+    for (uint32_t n = 0; n < N; ++n) {
+        for (uint32_t h = 0; h < H; ++h) {
+            uint32_t true_class = class_dist(gen);
+            target_tensor(n, h) = true_class;
+        }
+    }
+
+    auto input = core::from_xtensor(input_tensor, &autograd::ctx().get_device());
+    std::cout << "Input Logits:\n";
+    input.print();
+
+    auto target = core::from_xtensor<uint32_t, ttnn::DataType::UINT32>(
+        target_tensor, &autograd::ctx().get_device(), ttnn::Layout::ROW_MAJOR);
+    std::cout << "Input Target Indexes:\n";
+    target.print();
+
+    auto result = ttml::metal::cross_entropy_fw(input, target);
+    std::cout << "CrossEntropyForward_Test:\nResult:\n";
+    result.print();
+
+    auto expected_result = calculate_cross_entropy_loss(input_tensor, target_tensor);
+    auto expected_result_print = core::from_xtensor(expected_result, &autograd::ctx().get_device());
+    std::cout << "Expected Result:\n";
+    expected_result_print.print();
+
+    // Check if the result is close to the expected result
+    auto result_xtensor = core::to_xtensor(result);
+    assert((result_xtensor.shape() == expected_result.shape()));
+    EXPECT_TRUE(xt::allclose(result_xtensor, expected_result, 3e-2F, 1e-2F));
+}
+
+TEST_F(CrossEntropyForwardTest, CrossEntropyForward_Huge_Forward) {
+    using namespace ttml;
+
+    const uint32_t N = 64U, C = 1U, H = 32U, W = 128000U;
+    const auto shape = ttnn::SmallVector<size_t>{N, C, H, W};
+
+    std::random_device rd;
+    std::mt19937 gen(42);
+    xt::xarray<float> input_tensor = xt::random::rand<float>({N, C, H, W}, -10.0F, 10.0F, gen);
+    xt::xarray<uint32_t> target_tensor = xt::zeros<uint32_t>({N, H});
+
+    std::uniform_int_distribution<uint32_t> class_dist(0, W - 1);
+    for (uint32_t n = 0; n < N; ++n) {
+        for (uint32_t h = 0; h < H; ++h) {
+            uint32_t true_class = class_dist(gen);
+            target_tensor(n, h) = true_class;
+        }
+    }
+
+    auto input = core::from_xtensor(input_tensor, &autograd::ctx().get_device());
+    std::cout << "Input Logits:\n";
+    input.print();
+
+    auto target = core::from_xtensor<uint32_t, ttnn::DataType::UINT32>(
+        target_tensor, &autograd::ctx().get_device(), ttnn::Layout::ROW_MAJOR);
+    std::cout << "Input Target Indexes:\n";
+    target.print();
+
+    auto result = ttml::metal::cross_entropy_fw(input, target);
+    std::cout << "CrossEntropyForward_Test:\nResult:\n";
+    result.print();
+
+    auto expected_result = calculate_cross_entropy_loss(input_tensor, target_tensor);
+    auto expected_result_print = core::from_xtensor(expected_result, &autograd::ctx().get_device());
+    std::cout << "Expected Result:\n";
+    expected_result_print.print();
+
+    // Check if the result is close to the expected result
+    auto result_xtensor = core::to_xtensor(result);
+    assert((result_xtensor.shape() == expected_result.shape()));
+    EXPECT_TRUE(xt::allclose(result_xtensor, expected_result, 3e-2F, 1e-2F));
+}


Reviews:
user: login: rfurko-tt
body: Overall looks good. Great! 

Unfortunately, wasn't able to find bugs (why it does not work during first run after `tt-smi -r`)
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-12T07:05:26Z
user: login: vmelnykovTT
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-16T14:13:09Z
user: login: vmelnykovTT
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-16T14:25:35Z
user: login: vmelnykovTT
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-16T14:26:40Z
user: login: rfurko-tt
body: Looks good overall! Please clean it a bit more and ready to go 

@dmakoviichuk-tt could you please take a look as well? 
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-24T19:12:32Z
user: login: rfurko-tt
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-24T19:13:21Z
user: login: vmelnykovTT
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-25T13:07:49Z
user: login: vmelnykovTT
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-25T16:09:49Z
user: login: romaf5
state: COMMENTED
author_association: NONE
submitted_at: 2025-04-28T18:00:28Z
user: login: dmakoviichuk-tt
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-04-29T18:00:39Z

Review comments:
diff_hunk: @@ -0,0 +1,17 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include "cross_entropy_fw.hpp"
+
+#include "device/cross_entropy_fw_device_operation.hpp"
+
+namespace ttml::metal::ops::cross_entropy_fw {
+
+ttnn::Tensor CrossEntropyForwardOperation::invoke(const ttnn::Tensor& input_tensor, const ttnn::Tensor& target_tensor) {
+    auto result /*ttnn::Tensor*/ = ttnn::prim::ttml_cross_entropy_fw(
+        input_tensor, target_tensor);  // logits : model output (B, 1, S, D), target : ground truth (B, 1, S, D)
+    // resulut shape (B, 1, S, 1) - will change this later
path: tt-train/sources/ttml/metal/ops/cross_entropy_fw/cross_entropy_fw.cpp
user: login: rfurko-tt
body: remove this comment?
created_at: 2025-04-12T06:35:24Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,17 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include "cross_entropy_fw.hpp"
+
+#include "device/cross_entropy_fw_device_operation.hpp"
+
+namespace ttml::metal::ops::cross_entropy_fw {
+
+ttnn::Tensor CrossEntropyForwardOperation::invoke(const ttnn::Tensor& input_tensor, const ttnn::Tensor& target_tensor) {
+    auto result /*ttnn::Tensor*/ = ttnn::prim::ttml_cross_entropy_fw(
path: tt-train/sources/ttml/metal/ops/cross_entropy_fw/cross_entropy_fw.cpp
user: login: rfurko-tt
body: no need to write comment for `result` type 
created_at: 2025-04-12T06:35:35Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,17 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include "cross_entropy_fw.hpp"
+
+#include "device/cross_entropy_fw_device_operation.hpp"
+
+namespace ttml::metal::ops::cross_entropy_fw {
+
+ttnn::Tensor CrossEntropyForwardOperation::invoke(const ttnn::Tensor& input_tensor, const ttnn::Tensor& target_tensor) {
+    auto result /*ttnn::Tensor*/ = ttnn::prim::ttml_cross_entropy_fw(
path: tt-train/sources/ttml/metal/ops/cross_entropy_fw/cross_entropy_fw.cpp
user: login: rfurko-tt
body: if you need it, just change `auto` to `ttnn::Tensor`
created_at: 2025-04-12T06:35:50Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,132 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include "cross_entropy_fw_device_operation.hpp"
+
+#include "cross_entropy_fw_program_factory.hpp"
+
+namespace ttml::metal::ops::cross_entropy_fw::device {
+
+CrossEntropyForwardDeviceOperation::program_factory_t CrossEntropyForwardDeviceOperation::select_program_factory(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    return CrossEntropyForwardProgramFactory{};
+}
+
+void CrossEntropyForwardDeviceOperation::validate_on_program_cache_hit(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    validate_on_program_cache_miss(args, tensor_args);
+}
+
+void CrossEntropyForwardDeviceOperation::validate_on_program_cache_miss(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    // the same validation as rmsnorm
path: tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/cross_entropy_fw_device_operation.cpp
user: login: rfurko-tt
body: remove this comment
created_at: 2025-04-12T06:36:09Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,132 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include "cross_entropy_fw_device_operation.hpp"
+
+#include "cross_entropy_fw_program_factory.hpp"
+
+namespace ttml::metal::ops::cross_entropy_fw::device {
+
+CrossEntropyForwardDeviceOperation::program_factory_t CrossEntropyForwardDeviceOperation::select_program_factory(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    return CrossEntropyForwardProgramFactory{};
+}
+
+void CrossEntropyForwardDeviceOperation::validate_on_program_cache_hit(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    validate_on_program_cache_miss(args, tensor_args);
+}
+
+void CrossEntropyForwardDeviceOperation::validate_on_program_cache_miss(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    // the same validation as rmsnorm
+    auto check_tensor = [](const ttnn::Tensor& tensor, const std::string& name) {
+        TT_FATAL(
+            tensor.device()->arch() == tt::ARCH::WORMHOLE_B0,
+            "CrossEntropyForward operation is only supported on Wormhole. Device arch: {}. Tensor name {}",
+            magic_enum::enum_name(tensor.device()->arch()),
+            name);
+
+        TT_FATAL(
+            tensor.storage_type() == tt::tt_metal::StorageType::DEVICE,
+            "CrossEntropyForward operation requires {} to be on Device. Input storage type: {}",
+            name,
+            static_cast<int>(tensor.storage_type()));
+
+        TT_FATAL(
+            tensor.buffer() != nullptr,
+            "Operands to CrossEntropyForward need to be allocated in buffers on the device. Buffer is null. Tensor "
+            "name {}",
+            name);
+
+        TT_FATAL(
+            tensor.get_layout() == tt::tt_metal::Layout::TILE,
+            "CrossEntropyForward operation requires tensor to be in Tile layout. {} tensor layout: {}",
+            name,
+            static_cast<int>(tensor.get_layout()));
+
+        TT_FATAL(
+            tensor.get_dtype() == tt::tt_metal::DataType::BFLOAT16,
+            "CrossEntropyForward operation requires tensor to be of BFLOAT16 data type. {} tensor data type: {}",
+            name,
+            static_cast<int>(tensor.get_dtype()));
+
+        TT_FATAL(
+            tensor.memory_config().memory_layout == ttnn::TensorMemoryLayout::INTERLEAVED,
+            "CrossEntropyForward operation requires Interleaved memory layout. {} "
+            "memory layout: `{}`",
+            name,
+            static_cast<int>(tensor.memory_config().memory_layout));
+    };
+
+    const auto& input_tensor = tensor_args.input;
+    const auto& target_tensor = tensor_args.target;
+    const auto& preallocated_output_tensor = tensor_args.preallocated_output;
+    check_tensor(input_tensor, "Input");
+    check_tensor(target_tensor, "Target");
+    if (preallocated_output_tensor.has_value()) {
+        check_tensor(preallocated_output_tensor.value(), "Preallocated Output");
+    }
+}
+
+CrossEntropyForwardDeviceOperation::spec_return_value_t CrossEntropyForwardDeviceOperation::compute_output_specs(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    auto input_logical_shape = tensor_args.input.get_logical_shape();
+    // input_logical_shape[-1] = 1; // <- set the last dimension to 1
path: tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/cross_entropy_fw_device_operation.cpp
user: login: rfurko-tt
body: remove comment
created_at: 2025-04-12T06:36:20Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,132 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include "cross_entropy_fw_device_operation.hpp"
+
+#include "cross_entropy_fw_program_factory.hpp"
+
+namespace ttml::metal::ops::cross_entropy_fw::device {
+
+CrossEntropyForwardDeviceOperation::program_factory_t CrossEntropyForwardDeviceOperation::select_program_factory(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    return CrossEntropyForwardProgramFactory{};
+}
+
+void CrossEntropyForwardDeviceOperation::validate_on_program_cache_hit(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    validate_on_program_cache_miss(args, tensor_args);
+}
+
+void CrossEntropyForwardDeviceOperation::validate_on_program_cache_miss(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    // the same validation as rmsnorm
+    auto check_tensor = [](const ttnn::Tensor& tensor, const std::string& name) {
+        TT_FATAL(
+            tensor.device()->arch() == tt::ARCH::WORMHOLE_B0,
+            "CrossEntropyForward operation is only supported on Wormhole. Device arch: {}. Tensor name {}",
+            magic_enum::enum_name(tensor.device()->arch()),
+            name);
+
+        TT_FATAL(
+            tensor.storage_type() == tt::tt_metal::StorageType::DEVICE,
+            "CrossEntropyForward operation requires {} to be on Device. Input storage type: {}",
+            name,
+            static_cast<int>(tensor.storage_type()));
+
+        TT_FATAL(
+            tensor.buffer() != nullptr,
+            "Operands to CrossEntropyForward need to be allocated in buffers on the device. Buffer is null. Tensor "
+            "name {}",
+            name);
+
+        TT_FATAL(
+            tensor.get_layout() == tt::tt_metal::Layout::TILE,
+            "CrossEntropyForward operation requires tensor to be in Tile layout. {} tensor layout: {}",
+            name,
+            static_cast<int>(tensor.get_layout()));
+
+        TT_FATAL(
+            tensor.get_dtype() == tt::tt_metal::DataType::BFLOAT16,
+            "CrossEntropyForward operation requires tensor to be of BFLOAT16 data type. {} tensor data type: {}",
+            name,
+            static_cast<int>(tensor.get_dtype()));
+
+        TT_FATAL(
+            tensor.memory_config().memory_layout == ttnn::TensorMemoryLayout::INTERLEAVED,
+            "CrossEntropyForward operation requires Interleaved memory layout. {} "
+            "memory layout: `{}`",
+            name,
+            static_cast<int>(tensor.memory_config().memory_layout));
+    };
+
+    const auto& input_tensor = tensor_args.input;
+    const auto& target_tensor = tensor_args.target;
+    const auto& preallocated_output_tensor = tensor_args.preallocated_output;
+    check_tensor(input_tensor, "Input");
+    check_tensor(target_tensor, "Target");
+    if (preallocated_output_tensor.has_value()) {
+        check_tensor(preallocated_output_tensor.value(), "Preallocated Output");
+    }
+}
+
+CrossEntropyForwardDeviceOperation::spec_return_value_t CrossEntropyForwardDeviceOperation::compute_output_specs(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    auto input_logical_shape = tensor_args.input.get_logical_shape();
path: tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/cross_entropy_fw_device_operation.cpp
user: login: rfurko-tt
body: move this below if statement, why to request shape if we return existing tensor
created_at: 2025-04-12T06:36:45Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,132 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include "cross_entropy_fw_device_operation.hpp"
+
+#include "cross_entropy_fw_program_factory.hpp"
+
+namespace ttml::metal::ops::cross_entropy_fw::device {
+
+CrossEntropyForwardDeviceOperation::program_factory_t CrossEntropyForwardDeviceOperation::select_program_factory(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    return CrossEntropyForwardProgramFactory{};
+}
+
+void CrossEntropyForwardDeviceOperation::validate_on_program_cache_hit(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    validate_on_program_cache_miss(args, tensor_args);
+}
+
+void CrossEntropyForwardDeviceOperation::validate_on_program_cache_miss(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    // the same validation as rmsnorm
+    auto check_tensor = [](const ttnn::Tensor& tensor, const std::string& name) {
+        TT_FATAL(
+            tensor.device()->arch() == tt::ARCH::WORMHOLE_B0,
+            "CrossEntropyForward operation is only supported on Wormhole. Device arch: {}. Tensor name {}",
+            magic_enum::enum_name(tensor.device()->arch()),
+            name);
+
+        TT_FATAL(
+            tensor.storage_type() == tt::tt_metal::StorageType::DEVICE,
+            "CrossEntropyForward operation requires {} to be on Device. Input storage type: {}",
+            name,
+            static_cast<int>(tensor.storage_type()));
+
+        TT_FATAL(
+            tensor.buffer() != nullptr,
+            "Operands to CrossEntropyForward need to be allocated in buffers on the device. Buffer is null. Tensor "
+            "name {}",
+            name);
+
+        TT_FATAL(
+            tensor.get_layout() == tt::tt_metal::Layout::TILE,
+            "CrossEntropyForward operation requires tensor to be in Tile layout. {} tensor layout: {}",
+            name,
+            static_cast<int>(tensor.get_layout()));
+
+        TT_FATAL(
+            tensor.get_dtype() == tt::tt_metal::DataType::BFLOAT16,
+            "CrossEntropyForward operation requires tensor to be of BFLOAT16 data type. {} tensor data type: {}",
+            name,
+            static_cast<int>(tensor.get_dtype()));
+
+        TT_FATAL(
+            tensor.memory_config().memory_layout == ttnn::TensorMemoryLayout::INTERLEAVED,
+            "CrossEntropyForward operation requires Interleaved memory layout. {} "
+            "memory layout: `{}`",
+            name,
+            static_cast<int>(tensor.memory_config().memory_layout));
+    };
+
+    const auto& input_tensor = tensor_args.input;
+    const auto& target_tensor = tensor_args.target;
+    const auto& preallocated_output_tensor = tensor_args.preallocated_output;
+    check_tensor(input_tensor, "Input");
+    check_tensor(target_tensor, "Target");
+    if (preallocated_output_tensor.has_value()) {
+        check_tensor(preallocated_output_tensor.value(), "Preallocated Output");
+    }
+}
+
+CrossEntropyForwardDeviceOperation::spec_return_value_t CrossEntropyForwardDeviceOperation::compute_output_specs(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    auto input_logical_shape = tensor_args.input.get_logical_shape();
+    // input_logical_shape[-1] = 1; // <- set the last dimension to 1
+    if (tensor_args.preallocated_output.has_value()) {
+        return tensor_args.preallocated_output->get_tensor_spec();
+    }
+
+    input_logical_shape[-1] = 1;
+    return ttnn::TensorSpec(
+        ttnn::Shape(input_logical_shape),  // <- shape of the output tensor (B, 1, S, D) todo: change
path: tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/cross_entropy_fw_device_operation.cpp
user: login: rfurko-tt
body: remove comment 
created_at: 2025-04-12T06:36:55Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,132 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include "cross_entropy_fw_device_operation.hpp"
+
+#include "cross_entropy_fw_program_factory.hpp"
+
+namespace ttml::metal::ops::cross_entropy_fw::device {
+
+CrossEntropyForwardDeviceOperation::program_factory_t CrossEntropyForwardDeviceOperation::select_program_factory(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    return CrossEntropyForwardProgramFactory{};
+}
+
+void CrossEntropyForwardDeviceOperation::validate_on_program_cache_hit(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    validate_on_program_cache_miss(args, tensor_args);
+}
+
+void CrossEntropyForwardDeviceOperation::validate_on_program_cache_miss(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    // the same validation as rmsnorm
+    auto check_tensor = [](const ttnn::Tensor& tensor, const std::string& name) {
+        TT_FATAL(
+            tensor.device()->arch() == tt::ARCH::WORMHOLE_B0,
+            "CrossEntropyForward operation is only supported on Wormhole. Device arch: {}. Tensor name {}",
+            magic_enum::enum_name(tensor.device()->arch()),
+            name);
+
+        TT_FATAL(
+            tensor.storage_type() == tt::tt_metal::StorageType::DEVICE,
+            "CrossEntropyForward operation requires {} to be on Device. Input storage type: {}",
+            name,
+            static_cast<int>(tensor.storage_type()));
+
+        TT_FATAL(
+            tensor.buffer() != nullptr,
+            "Operands to CrossEntropyForward need to be allocated in buffers on the device. Buffer is null. Tensor "
+            "name {}",
+            name);
+
+        TT_FATAL(
+            tensor.get_layout() == tt::tt_metal::Layout::TILE,
+            "CrossEntropyForward operation requires tensor to be in Tile layout. {} tensor layout: {}",
+            name,
+            static_cast<int>(tensor.get_layout()));
+
+        TT_FATAL(
+            tensor.get_dtype() == tt::tt_metal::DataType::BFLOAT16,
+            "CrossEntropyForward operation requires tensor to be of BFLOAT16 data type. {} tensor data type: {}",
+            name,
+            static_cast<int>(tensor.get_dtype()));
+
+        TT_FATAL(
+            tensor.memory_config().memory_layout == ttnn::TensorMemoryLayout::INTERLEAVED,
+            "CrossEntropyForward operation requires Interleaved memory layout. {} "
+            "memory layout: `{}`",
+            name,
+            static_cast<int>(tensor.memory_config().memory_layout));
+    };
+
+    const auto& input_tensor = tensor_args.input;
+    const auto& target_tensor = tensor_args.target;
+    const auto& preallocated_output_tensor = tensor_args.preallocated_output;
+    check_tensor(input_tensor, "Input");
+    check_tensor(target_tensor, "Target");
+    if (preallocated_output_tensor.has_value()) {
+        check_tensor(preallocated_output_tensor.value(), "Preallocated Output");
+    }
+}
+
+CrossEntropyForwardDeviceOperation::spec_return_value_t CrossEntropyForwardDeviceOperation::compute_output_specs(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    auto input_logical_shape = tensor_args.input.get_logical_shape();
+    // input_logical_shape[-1] = 1; // <- set the last dimension to 1
+    if (tensor_args.preallocated_output.has_value()) {
+        return tensor_args.preallocated_output->get_tensor_spec();
+    }
+
+    input_logical_shape[-1] = 1;
+    return ttnn::TensorSpec(
+        ttnn::Shape(input_logical_shape),  // <- shape of the output tensor (B, 1, S, D) todo: change
+                                           // to (B, 1, S, 1)
+        tt::tt_metal::TensorLayout(
+            tensor_args.input.get_dtype(), tt::tt_metal::Layout::TILE, tensor_args.input.memory_config()));
+}
+
+CrossEntropyForwardDeviceOperation::tensor_return_value_t CrossEntropyForwardDeviceOperation::create_output_tensors(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    tensor_return_value_t output_tensor;  // <- ttnn::Tensor
+
+    spec_return_value_t output_specs = compute_output_specs(args, tensor_args);
+
+    if (tensor_args.preallocated_output.has_value()) {
+        output_tensor = tensor_args.preallocated_output.value();
+    } else {
+        output_tensor = create_device_tensor(output_specs, tensor_args.input.device());
+    }
+
+    return output_tensor;
+}
+
+tt::stl::hash::hash_t CrossEntropyForwardDeviceOperation::compute_program_hash(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    // to do: compute program hash based on input tensor shape
path: tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/cross_entropy_fw_device_operation.cpp
user: login: rfurko-tt
body: remove todo
created_at: 2025-04-12T06:37:16Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,38 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include "dataflow_api.h"
+
+void kernel_main() {
+    uint32_t runtime_args_counter = 0;
+    uint32_t output_addr = get_arg_val<uint32_t>(runtime_args_counter++);
+    uint32_t num_rows_to_process = get_arg_val<uint32_t>(runtime_args_counter++);
+    uint32_t start_row = get_arg_val<uint32_t>(runtime_args_counter++);
+
+    constexpr uint32_t cb_output_idx = tt::CBIndex::c_10;
+
+    constexpr uint32_t block_size = get_compile_time_arg_val(0);  // size of block == 1 tile
path: tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/kernels/dataflow/writer_cross_entropy_fw_interleaved_start_id.cpp
user: login: rfurko-tt
body: comment is not correct
created_at: 2025-04-12T06:39:20Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,103 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include <cstdint>
+#include <cstring>
+
+#include "dataflow_api.h"
+#include "debug/dprint.h"
+#include "debug/dprint_pages.h"
+
+void kernel_main() {
+    uint32_t runtime_args_counter = 0U;
+    uint32_t input_address = get_arg_val<uint32_t>(runtime_args_counter++);        // input buffer address
+    uint32_t target_address = get_arg_val<uint32_t>(runtime_args_counter++);       // target buffer address
+    uint32_t num_rows_to_process = get_arg_val<uint32_t>(runtime_args_counter++);  // rows to process in this kernel
+    uint32_t start_row =
+        get_arg_val<uint32_t>(runtime_args_counter++);  // pre calculated num_rows_written in program factory
+
+    constexpr uint32_t cb_input_idx = tt::CBIndex::c_0;
+    constexpr uint32_t cb_target_idx = tt::CBIndex::c_1;
+    constexpr uint32_t cb_mask_idx = tt::CBIndex::c_2;
+    constexpr uint32_t cb_max_mask_idx = tt::CBIndex::c_3;
+    constexpr uint32_t cb_scaler_idx = tt::CBIndex::c_4;
+
+    constexpr uint32_t num_input_tiles = get_compile_time_arg_val(0);
path: tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/kernels/dataflow/reader_cross_entropy_fw_interleaved_start_id.cpp
user: login: rfurko-tt
body: remove `num_input_tiles` as it is not used
created_at: 2025-04-12T06:41:36Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,474 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include "cross_entropy_fw_program_factory.hpp"
+
+#include <bit>
+#include <cstdint>
+
+#include "cross_entropy_fw_device_operation_types.hpp"
+
+namespace {
+
+constexpr auto kReaderKernelPath =
+    "tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/kernels/dataflow/"
+    "reader_cross_entropy_fw_interleaved_start_id.cpp";
+
+constexpr auto kWriterKernelPath =
+    "tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/kernels/dataflow/"
+    "writer_cross_entropy_fw_interleaved_start_id.cpp";
+
+constexpr auto kComputeKernelPath =
+    "tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/kernels/compute/cross_entropy_fw_kernel.cpp";
+
+// reader runtime args
+constexpr uint32_t kInputBufferIdx = 0;
+constexpr uint32_t kTargetBufferIdx = 1U;
+// writer runtime args
+constexpr uint32_t kOutputBufferIdx = 0;
+
+constexpr auto kInputCbIndex = tt::CBIndex::c_0;
+constexpr auto kTargetCbIndex = tt::CBIndex::c_1;
+constexpr auto kMaskCbIndex = tt::CBIndex::c_2;
+constexpr auto kMaxMaskCbIndex = tt::CBIndex::c_3;
+constexpr auto kScalerCbIndex = tt::CBIndex::c_4;  // used to reduction
path: tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/cross_entropy_fw_program_factory.cpp
user: login: rfurko-tt
body: for reduction
created_at: 2025-04-12T06:52:14Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,474 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include "cross_entropy_fw_program_factory.hpp"
+
+#include <bit>
+#include <cstdint>
+
+#include "cross_entropy_fw_device_operation_types.hpp"
+
+namespace {
+
+constexpr auto kReaderKernelPath =
+    "tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/kernels/dataflow/"
+    "reader_cross_entropy_fw_interleaved_start_id.cpp";
+
+constexpr auto kWriterKernelPath =
+    "tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/kernels/dataflow/"
+    "writer_cross_entropy_fw_interleaved_start_id.cpp";
+
+constexpr auto kComputeKernelPath =
+    "tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/kernels/compute/cross_entropy_fw_kernel.cpp";
+
+// reader runtime args
+constexpr uint32_t kInputBufferIdx = 0;
+constexpr uint32_t kTargetBufferIdx = 1U;
+// writer runtime args
+constexpr uint32_t kOutputBufferIdx = 0;
+
+constexpr auto kInputCbIndex = tt::CBIndex::c_0;
+constexpr auto kTargetCbIndex = tt::CBIndex::c_1;
+constexpr auto kMaskCbIndex = tt::CBIndex::c_2;
+constexpr auto kMaxMaskCbIndex = tt::CBIndex::c_3;
+constexpr auto kScalerCbIndex = tt::CBIndex::c_4;  // used to reduction
+constexpr auto kMaxValueBeforeReductionCbIndex = tt::CBIndex::c_5;
+constexpr auto kMaxValueAfterReductionCbIndex = tt::CBIndex::c_6;
+constexpr auto kExpSumBeforeReductionCbIndex = tt::CBIndex::c_7;
+constexpr auto KExpSumAfterReductionCbIndex = tt::CBIndex::c_8;
+constexpr auto kOutputBeforeReductionCbIndex = tt::CBIndex::c_9;
+constexpr auto kOutputCbIndex = tt::CBIndex::c_10;
+
+constexpr uint32_t kNumMaskTiles = 1U;
+constexpr uint32_t kNumMaxValueTiles = 1U;
+constexpr uint32_t kNumOutputTiles = 1U;
+constexpr uint32_t kMaxValueBeforeReductionTiles = 1U;
+constexpr uint32_t kNumMaxValueAfterReductionTiles = 1U;
+constexpr uint32_t kNumExpSumBeforeReductionTiles = 1U;
+constexpr uint32_t kNumExpSumAfterReductionTiles = 1U;
+constexpr uint32_t kNumScalerTiles = 1U;  // used it to reduction ???
+constexpr uint32_t kNumOutputBeforeReductionTiles = 1U;
+
+const std::string kMaskWDefineKey = "DO_MASK_W";
+
+uint32_t get_block_size(uint32_t num_inner) {
+    const uint32_t max_block_size = 4U;  // 4 is the maximum block size for enabled fp32 dest acc
+    for (uint32_t block_size = max_block_size; block_size > 1; block_size--) {
+        if (num_inner % block_size == 0) {  // if num_inner is divisible by block_size - choose this block_size
+            return block_size;
+        }
+    }
+    return 1U;
+}
+
+}  // namespace
+
+namespace ttml::metal::ops::cross_entropy_fw::device {
+
+/**
+ *   Helper struct to hold references to all kernels we create,
+ *        used during runtime argument setup.
+ */
+struct CrossEntropyForwardKernels {
+    tt::tt_metal::KernelHandle reader;
+    tt::tt_metal::KernelHandle writer;
+    tt::tt_metal::KernelHandle compute_group_1;
+    tt::tt_metal::KernelHandle compute_group_2;
+};
+
+/**
+ *   Create and configure a circular buffer, returning both the configuration and the handle.
+ */
+tt::tt_metal::CBHandle create_circular_buffer(
+    tt::tt_metal::Program& program,
+    const tt::tt_metal::CoreRangeSet& core_ranges,
+    uint32_t cb_index,
+    tt::DataFormat data_format,
+    uint32_t single_tile_size,
+    uint32_t num_tiles) {
+    tt::tt_metal::CircularBufferConfig cb_config =
+        tt::tt_metal::CircularBufferConfig(num_tiles * single_tile_size, {{cb_index, data_format}})
+            .set_page_size(cb_index, single_tile_size);
+
+    auto cb_handle = CreateCircularBuffer(program, core_ranges, cb_config);
+    return cb_handle;
+}
+
+/**
+ *   Create a reader kernel with the given compile-time arguments.
+ */
+tt::tt_metal::KernelHandle create_reader_kernel(
+    tt::tt_metal::Program& program,
+    const tt::tt_metal::CoreRangeSet& core_ranges,
+    const std::vector<uint32_t>& compile_time_args,
+    const std::map<std::string, std::string>& defines,
+    const std::string& kernel_path) {
+    return tt::tt_metal::CreateKernel(
+        program, kernel_path, core_ranges, tt::tt_metal::ReaderDataMovementConfig(compile_time_args, defines));
+}
+
+/**
+ *   Create a writer kernel with the given compile-time arguments.
+ */
+tt::tt_metal::KernelHandle create_writer_kernel(
+    tt::tt_metal::Program& program,
+    const tt::tt_metal::CoreRangeSet& core_ranges,
+    const std::vector<uint32_t>& compile_time_args,
+    const std::map<std::string, std::string>& defines,
+    const std::string& kernel_path) {
+    return tt::tt_metal::CreateKernel(
+        program, kernel_path, core_ranges, tt::tt_metal::WriterDataMovementConfig(compile_time_args, defines));
+}
+
+/**
+ * Create a compute kernel with the given compile-time arguments.
+ */
+tt::tt_metal::KernelHandle create_compute_kernel(
+    tt::tt_metal::Program& program,
+    const tt::tt_metal::CoreRangeSet& core_ranges,
+    const std::vector<uint32_t>& compile_time_args,
+    const std::map<std::string, std::string>& defines,
+    const std::string& kernel_path) {
+    return tt::tt_metal::CreateKernel(
+        program,
+        kernel_path,
+        core_ranges,
+        tt::tt_metal::ComputeConfig{
+            .math_fidelity = MathFidelity::HiFi4,
+            .fp32_dest_acc_en = true,
+            .math_approx_mode = false,
+            .compile_args = compile_time_args,
+            .defines = defines});
+}
+
+/**
+ * Set up the runtime arguments for the 4 relevant kernels (reader, writer, compute G1, compute G2)
+ *        for each core in the grid.
+ */
+void assign_per_core_runtime_args(
+    tt::tt_metal::Program& program,
+    const CrossEntropyForwardKernels& kernels,
+    const tt::tt_metal::Buffer* input_buffer,
+    const tt::tt_metal::Buffer* target_buffer,
+    const tt::tt_metal::Buffer* output_buffer,
+    uint32_t num_cores,
+    uint32_t num_cores_y,
+    uint32_t num_rows_per_core_group_1,
+    uint32_t num_rows_per_core_group_2,
+    const tt::tt_metal::CoreRangeSet& core_group_1,
+    const tt::tt_metal::CoreRangeSet& core_group_2) {
+    for (uint32_t i = 0, num_rows_written = 0; i < num_cores; i++) {
+        CoreCoord core = {i / num_cores_y, i % num_cores_y};
+
+        // Determine how many rows this core will process
+        uint32_t num_rows_per_core = 0;
+        if (core_group_1.contains(core)) {
+            num_rows_per_core = num_rows_per_core_group_1;
+        } else if (core_group_2.contains(core)) {
+            num_rows_per_core = num_rows_per_core_group_2;
+        } else {
+            TT_FATAL(false, "Core not in specified core ranges");
+        }
+
+        // Reader kernel: (input_addr, gamma_addr, number_of_rows, offset_in_rows)
+        SetRuntimeArgs(
+            program,
+            kernels.reader,
+            core,
+            {input_buffer->address(), target_buffer->address(), num_rows_per_core, num_rows_written});
+
+        // Writer kernel: (dst_addr, dst_rms_addr number_of_rows, offset_in_rows)
+        SetRuntimeArgs(
+            program,
+            kernels.writer,
+            core,
+            {output_buffer->address(), /*temporary*/ num_rows_per_core, num_rows_written});
+
+        num_rows_written += num_rows_per_core;
+    }
+}
+
+CrossEntropyForwardProgramFactory::cached_program_t CrossEntropyForwardProgramFactory::create(
+    const operation_attributes_t& operation_attributes,
+    const tensor_args_t& tensor_args,
+    tensor_return_value_t& output) {
+    // -------------------------------------------------------------------------
+    // 1) Setup device, data formats, tile sizes, and compute split
+    // -------------------------------------------------------------------------
+    const auto& input = tensor_args.input;
+    const auto& target = tensor_args.target;
+
+    auto* device = input.device();
+
+    tt::tt_metal::Program program{};
+
+    tt::DataFormat input_data_format = datatype_to_dataformat_converter(input.get_dtype());
+    TT_FATAL(input_data_format == tt::DataFormat::Float16_b, "Input data format must be Float16_b");
+
+    uint32_t bfloat16_single_tile_size_bytes = tt::tt_metal::detail::TileSize(tt::DataFormat::Float16_b);
+    uint32_t float32_single_tile_size_bytes = tt::tt_metal::detail::TileSize(tt::DataFormat::Float32);
+
+    auto padded_tensor_shape = input.get_padded_shape();
+    auto padded_tensor_volume = input.volume();
+
+    TT_FATAL(
+        padded_tensor_volume % tt::constants::TILE_HW == 0, "Padded input tensor volume must be divisible by TILE_HW");
+    TT_FATAL(padded_tensor_shape.rank() == 4U, "Input tensor must be 4D");
+
+    //
+    uint32_t Wt = padded_tensor_shape[-1] / tt::constants::TILE_WIDTH;  // <- number of tiles in inner dimension
+    uint32_t Ht = padded_tensor_shape[-2] / tt::constants::TILE_HEIGHT;
+    uint32_t NC = padded_tensor_shape[0] * padded_tensor_shape[1];
+    uint32_t total_rows_to_process = NC * Ht;
+
+    // get number of free cores
+    auto compute_with_storage_grid_size = device->compute_with_storage_grid_size();
+    uint32_t num_cores_x = compute_with_storage_grid_size.x;
+    uint32_t num_cores_y = compute_with_storage_grid_size.y;
+
+    // get the number of inner dimension(D) - we will caclulate log softmax on this dimension
+    uint32_t num_inner = input.get_logical_shape()[-1];  // (B, 1, S, D) and we take D here
+
+    // mask_w - this mask used to avoid calculation of extra data(data which will be added to create full tile 32x32)??
+    uint32_t mask_w = num_inner % tt::constants::TILE_WIDTH;  // <- width index of first trash value in tile
+
+    // compile arguments
+    uint32_t block_size = get_block_size(Wt);
+
+    auto [num_cores, all_cores, core_group_1, core_group_2, num_rows_per_core_group_1, num_rows_per_core_group_2] =
+        tt::tt_metal::split_work_to_cores(compute_with_storage_grid_size, total_rows_to_process);
+
+    // -------------------------------------------------------------------------
+    // 2) Create and configure circular buffers
+    // -------------------------------------------------------------------------
+
+    // TODO: add calculation of required  memory and check if it fits in L1
+    // then add possibility to kernels to process data by blocks
+    const uint32_t available_L1_in_bytes =
+        device->l1_size_per_core() - device->allocator()->get_base_allocator_addr(tt::tt_metal::HalMemType::L1);
+
+    const uint64_t mask_scaler_maxmask_memory =
+        (kNumMaskTiles + kNumScalerTiles + kNumMaskTiles) * bfloat16_single_tile_size_bytes;
+
+    const uint32_t num_input_tiles = Wt;
+    const uint32_t num_target_tiles = Wt;
+
+    auto data_format = input_data_format;  // tt::DataFormat::Float16_b
+    auto precise_data_format = tt::DataFormat::Float32;
+
+    auto cb_input = create_circular_buffer(
+        program, all_cores, kInputCbIndex, data_format, bfloat16_single_tile_size_bytes, num_input_tiles);
+
+    auto cb_target = create_circular_buffer(
+        program, all_cores, kTargetCbIndex, data_format, bfloat16_single_tile_size_bytes, num_target_tiles);
+
+    auto cb_mask = create_circular_buffer(
+        program, all_cores, kMaskCbIndex, data_format, bfloat16_single_tile_size_bytes, kNumMaskTiles);
+
+    auto cb_max_mask = create_circular_buffer(
+        program, all_cores, kMaxMaskCbIndex, data_format, bfloat16_single_tile_size_bytes, kNumMaskTiles);
+
+    auto cb_scaler = create_circular_buffer(
+        program, all_cores, kScalerCbIndex, data_format, bfloat16_single_tile_size_bytes, kNumScalerTiles);
+
+    auto cb_max_value_before_reduction = create_circular_buffer(
+        program,
+        all_cores,
+        kMaxValueBeforeReductionCbIndex,
+        data_format,
+        bfloat16_single_tile_size_bytes,
+        kMaxValueBeforeReductionTiles);
+
+    auto cb_max_value_after_reduction = create_circular_buffer(
+        program,
+        all_cores,
+        kMaxValueAfterReductionCbIndex,
+        data_format,
+        bfloat16_single_tile_size_bytes,
+        kNumMaxValueAfterReductionTiles);
+
+    auto cb_exp_sum_before_reduction = create_circular_buffer(
+        program,
+        all_cores,
+        kExpSumBeforeReductionCbIndex,
+        data_format,
+        bfloat16_single_tile_size_bytes,
+        kNumExpSumBeforeReductionTiles);
+
+    auto cb_exp_sum_after_refuction = create_circular_buffer(
+        program,
+        all_cores,
+        KExpSumAfterReductionCbIndex,
+        data_format,
+        bfloat16_single_tile_size_bytes,
+        kNumExpSumAfterReductionTiles);
+
+    auto cb_output_before_refuction = create_circular_buffer(
+        program,
+        all_cores,
+        kOutputBeforeReductionCbIndex,
+        data_format,
+        bfloat16_single_tile_size_bytes,
+        kNumOutputBeforeReductionTiles);
+
+    auto cb_output = create_circular_buffer(
+        program,
+        all_cores,
+        kOutputCbIndex,
+        data_format,
+        bfloat16_single_tile_size_bytes,
+        2 * block_size);  // create twice block_size variable
path: tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/cross_entropy_fw_program_factory.cpp
user: login: rfurko-tt
body: substitute it to just 2 (we don't need block here)
created_at: 2025-04-12T06:54:15Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,474 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include "cross_entropy_fw_program_factory.hpp"
+
+#include <bit>
+#include <cstdint>
+
+#include "cross_entropy_fw_device_operation_types.hpp"
+
+namespace {
+
+constexpr auto kReaderKernelPath =
+    "tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/kernels/dataflow/"
+    "reader_cross_entropy_fw_interleaved_start_id.cpp";
+
+constexpr auto kWriterKernelPath =
+    "tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/kernels/dataflow/"
+    "writer_cross_entropy_fw_interleaved_start_id.cpp";
+
+constexpr auto kComputeKernelPath =
+    "tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/kernels/compute/cross_entropy_fw_kernel.cpp";
+
+// reader runtime args
+constexpr uint32_t kInputBufferIdx = 0;
+constexpr uint32_t kTargetBufferIdx = 1U;
+// writer runtime args
+constexpr uint32_t kOutputBufferIdx = 0;
+
+constexpr auto kInputCbIndex = tt::CBIndex::c_0;
+constexpr auto kTargetCbIndex = tt::CBIndex::c_1;
+constexpr auto kMaskCbIndex = tt::CBIndex::c_2;
+constexpr auto kMaxMaskCbIndex = tt::CBIndex::c_3;
+constexpr auto kScalerCbIndex = tt::CBIndex::c_4;  // used to reduction
+constexpr auto kMaxValueBeforeReductionCbIndex = tt::CBIndex::c_5;
+constexpr auto kMaxValueAfterReductionCbIndex = tt::CBIndex::c_6;
+constexpr auto kExpSumBeforeReductionCbIndex = tt::CBIndex::c_7;
+constexpr auto KExpSumAfterReductionCbIndex = tt::CBIndex::c_8;
+constexpr auto kOutputBeforeReductionCbIndex = tt::CBIndex::c_9;
+constexpr auto kOutputCbIndex = tt::CBIndex::c_10;
+
+constexpr uint32_t kNumMaskTiles = 1U;
+constexpr uint32_t kNumMaxValueTiles = 1U;
+constexpr uint32_t kNumOutputTiles = 1U;
+constexpr uint32_t kMaxValueBeforeReductionTiles = 1U;
+constexpr uint32_t kNumMaxValueAfterReductionTiles = 1U;
+constexpr uint32_t kNumExpSumBeforeReductionTiles = 1U;
+constexpr uint32_t kNumExpSumAfterReductionTiles = 1U;
+constexpr uint32_t kNumScalerTiles = 1U;  // used it to reduction ???
+constexpr uint32_t kNumOutputBeforeReductionTiles = 1U;
+
+const std::string kMaskWDefineKey = "DO_MASK_W";
+
+uint32_t get_block_size(uint32_t num_inner) {
+    const uint32_t max_block_size = 4U;  // 4 is the maximum block size for enabled fp32 dest acc
+    for (uint32_t block_size = max_block_size; block_size > 1; block_size--) {
+        if (num_inner % block_size == 0) {  // if num_inner is divisible by block_size - choose this block_size
+            return block_size;
+        }
+    }
+    return 1U;
+}
+
+}  // namespace
+
+namespace ttml::metal::ops::cross_entropy_fw::device {
+
+/**
+ *   Helper struct to hold references to all kernels we create,
+ *        used during runtime argument setup.
+ */
+struct CrossEntropyForwardKernels {
+    tt::tt_metal::KernelHandle reader;
+    tt::tt_metal::KernelHandle writer;
+    tt::tt_metal::KernelHandle compute_group_1;
+    tt::tt_metal::KernelHandle compute_group_2;
+};
+
+/**
+ *   Create and configure a circular buffer, returning both the configuration and the handle.
+ */
+tt::tt_metal::CBHandle create_circular_buffer(
+    tt::tt_metal::Program& program,
+    const tt::tt_metal::CoreRangeSet& core_ranges,
+    uint32_t cb_index,
+    tt::DataFormat data_format,
+    uint32_t single_tile_size,
+    uint32_t num_tiles) {
+    tt::tt_metal::CircularBufferConfig cb_config =
+        tt::tt_metal::CircularBufferConfig(num_tiles * single_tile_size, {{cb_index, data_format}})
+            .set_page_size(cb_index, single_tile_size);
+
+    auto cb_handle = CreateCircularBuffer(program, core_ranges, cb_config);
+    return cb_handle;
+}
+
+/**
+ *   Create a reader kernel with the given compile-time arguments.
+ */
+tt::tt_metal::KernelHandle create_reader_kernel(
+    tt::tt_metal::Program& program,
+    const tt::tt_metal::CoreRangeSet& core_ranges,
+    const std::vector<uint32_t>& compile_time_args,
+    const std::map<std::string, std::string>& defines,
+    const std::string& kernel_path) {
+    return tt::tt_metal::CreateKernel(
+        program, kernel_path, core_ranges, tt::tt_metal::ReaderDataMovementConfig(compile_time_args, defines));
+}
+
+/**
+ *   Create a writer kernel with the given compile-time arguments.
+ */
+tt::tt_metal::KernelHandle create_writer_kernel(
+    tt::tt_metal::Program& program,
+    const tt::tt_metal::CoreRangeSet& core_ranges,
+    const std::vector<uint32_t>& compile_time_args,
+    const std::map<std::string, std::string>& defines,
+    const std::string& kernel_path) {
+    return tt::tt_metal::CreateKernel(
+        program, kernel_path, core_ranges, tt::tt_metal::WriterDataMovementConfig(compile_time_args, defines));
+}
+
+/**
+ * Create a compute kernel with the given compile-time arguments.
+ */
+tt::tt_metal::KernelHandle create_compute_kernel(
+    tt::tt_metal::Program& program,
+    const tt::tt_metal::CoreRangeSet& core_ranges,
+    const std::vector<uint32_t>& compile_time_args,
+    const std::map<std::string, std::string>& defines,
+    const std::string& kernel_path) {
+    return tt::tt_metal::CreateKernel(
+        program,
+        kernel_path,
+        core_ranges,
+        tt::tt_metal::ComputeConfig{
+            .math_fidelity = MathFidelity::HiFi4,
+            .fp32_dest_acc_en = true,
+            .math_approx_mode = false,
+            .compile_args = compile_time_args,
+            .defines = defines});
+}
+
+/**
+ * Set up the runtime arguments for the 4 relevant kernels (reader, writer, compute G1, compute G2)
+ *        for each core in the grid.
+ */
+void assign_per_core_runtime_args(
+    tt::tt_metal::Program& program,
+    const CrossEntropyForwardKernels& kernels,
+    const tt::tt_metal::Buffer* input_buffer,
+    const tt::tt_metal::Buffer* target_buffer,
+    const tt::tt_metal::Buffer* output_buffer,
+    uint32_t num_cores,
+    uint32_t num_cores_y,
+    uint32_t num_rows_per_core_group_1,
+    uint32_t num_rows_per_core_group_2,
+    const tt::tt_metal::CoreRangeSet& core_group_1,
+    const tt::tt_metal::CoreRangeSet& core_group_2) {
+    for (uint32_t i = 0, num_rows_written = 0; i < num_cores; i++) {
+        CoreCoord core = {i / num_cores_y, i % num_cores_y};
+
+        // Determine how many rows this core will process
+        uint32_t num_rows_per_core = 0;
+        if (core_group_1.contains(core)) {
+            num_rows_per_core = num_rows_per_core_group_1;
+        } else if (core_group_2.contains(core)) {
+            num_rows_per_core = num_rows_per_core_group_2;
+        } else {
+            TT_FATAL(false, "Core not in specified core ranges");
+        }
+
+        // Reader kernel: (input_addr, gamma_addr, number_of_rows, offset_in_rows)
+        SetRuntimeArgs(
+            program,
+            kernels.reader,
+            core,
+            {input_buffer->address(), target_buffer->address(), num_rows_per_core, num_rows_written});
+
+        // Writer kernel: (dst_addr, dst_rms_addr number_of_rows, offset_in_rows)
+        SetRuntimeArgs(
+            program,
+            kernels.writer,
+            core,
+            {output_buffer->address(), /*temporary*/ num_rows_per_core, num_rows_written});
+
+        num_rows_written += num_rows_per_core;
+    }
+}
+
+CrossEntropyForwardProgramFactory::cached_program_t CrossEntropyForwardProgramFactory::create(
+    const operation_attributes_t& operation_attributes,
+    const tensor_args_t& tensor_args,
+    tensor_return_value_t& output) {
+    // -------------------------------------------------------------------------
+    // 1) Setup device, data formats, tile sizes, and compute split
+    // -------------------------------------------------------------------------
+    const auto& input = tensor_args.input;
+    const auto& target = tensor_args.target;
+
+    auto* device = input.device();
+
+    tt::tt_metal::Program program{};
+
+    tt::DataFormat input_data_format = datatype_to_dataformat_converter(input.get_dtype());
+    TT_FATAL(input_data_format == tt::DataFormat::Float16_b, "Input data format must be Float16_b");
+
+    uint32_t bfloat16_single_tile_size_bytes = tt::tt_metal::detail::TileSize(tt::DataFormat::Float16_b);
+    uint32_t float32_single_tile_size_bytes = tt::tt_metal::detail::TileSize(tt::DataFormat::Float32);
+
+    auto padded_tensor_shape = input.get_padded_shape();
+    auto padded_tensor_volume = input.volume();
+
+    TT_FATAL(
+        padded_tensor_volume % tt::constants::TILE_HW == 0, "Padded input tensor volume must be divisible by TILE_HW");
+    TT_FATAL(padded_tensor_shape.rank() == 4U, "Input tensor must be 4D");
+
+    //
+    uint32_t Wt = padded_tensor_shape[-1] / tt::constants::TILE_WIDTH;  // <- number of tiles in inner dimension
+    uint32_t Ht = padded_tensor_shape[-2] / tt::constants::TILE_HEIGHT;
+    uint32_t NC = padded_tensor_shape[0] * padded_tensor_shape[1];
+    uint32_t total_rows_to_process = NC * Ht;
+
+    // get number of free cores
+    auto compute_with_storage_grid_size = device->compute_with_storage_grid_size();
+    uint32_t num_cores_x = compute_with_storage_grid_size.x;
+    uint32_t num_cores_y = compute_with_storage_grid_size.y;
+
+    // get the number of inner dimension(D) - we will caclulate log softmax on this dimension
+    uint32_t num_inner = input.get_logical_shape()[-1];  // (B, 1, S, D) and we take D here
+
+    // mask_w - this mask used to avoid calculation of extra data(data which will be added to create full tile 32x32)??
+    uint32_t mask_w = num_inner % tt::constants::TILE_WIDTH;  // <- width index of first trash value in tile
+
+    // compile arguments
+    uint32_t block_size = get_block_size(Wt);
+
+    auto [num_cores, all_cores, core_group_1, core_group_2, num_rows_per_core_group_1, num_rows_per_core_group_2] =
+        tt::tt_metal::split_work_to_cores(compute_with_storage_grid_size, total_rows_to_process);
+
+    // -------------------------------------------------------------------------
+    // 2) Create and configure circular buffers
+    // -------------------------------------------------------------------------
+
+    // TODO: add calculation of required  memory and check if it fits in L1
+    // then add possibility to kernels to process data by blocks
+    const uint32_t available_L1_in_bytes =
+        device->l1_size_per_core() - device->allocator()->get_base_allocator_addr(tt::tt_metal::HalMemType::L1);
+
+    const uint64_t mask_scaler_maxmask_memory =
+        (kNumMaskTiles + kNumScalerTiles + kNumMaskTiles) * bfloat16_single_tile_size_bytes;
+
+    const uint32_t num_input_tiles = Wt;
+    const uint32_t num_target_tiles = Wt;
+
+    auto data_format = input_data_format;  // tt::DataFormat::Float16_b
+    auto precise_data_format = tt::DataFormat::Float32;
+
+    auto cb_input = create_circular_buffer(
+        program, all_cores, kInputCbIndex, data_format, bfloat16_single_tile_size_bytes, num_input_tiles);
+
+    auto cb_target = create_circular_buffer(
+        program, all_cores, kTargetCbIndex, data_format, bfloat16_single_tile_size_bytes, num_target_tiles);
+
+    auto cb_mask = create_circular_buffer(
+        program, all_cores, kMaskCbIndex, data_format, bfloat16_single_tile_size_bytes, kNumMaskTiles);
+
+    auto cb_max_mask = create_circular_buffer(
+        program, all_cores, kMaxMaskCbIndex, data_format, bfloat16_single_tile_size_bytes, kNumMaskTiles);
+
+    auto cb_scaler = create_circular_buffer(
+        program, all_cores, kScalerCbIndex, data_format, bfloat16_single_tile_size_bytes, kNumScalerTiles);
+
+    auto cb_max_value_before_reduction = create_circular_buffer(
+        program,
+        all_cores,
+        kMaxValueBeforeReductionCbIndex,
+        data_format,
+        bfloat16_single_tile_size_bytes,
+        kMaxValueBeforeReductionTiles);
+
+    auto cb_max_value_after_reduction = create_circular_buffer(
+        program,
+        all_cores,
+        kMaxValueAfterReductionCbIndex,
+        data_format,
+        bfloat16_single_tile_size_bytes,
+        kNumMaxValueAfterReductionTiles);
+
+    auto cb_exp_sum_before_reduction = create_circular_buffer(
+        program,
+        all_cores,
+        kExpSumBeforeReductionCbIndex,
+        data_format,
+        bfloat16_single_tile_size_bytes,
+        kNumExpSumBeforeReductionTiles);
+
+    auto cb_exp_sum_after_refuction = create_circular_buffer(
+        program,
+        all_cores,
+        KExpSumAfterReductionCbIndex,
+        data_format,
+        bfloat16_single_tile_size_bytes,
+        kNumExpSumAfterReductionTiles);
+
+    auto cb_output_before_refuction = create_circular_buffer(
+        program,
+        all_cores,
+        kOutputBeforeReductionCbIndex,
+        data_format,
+        bfloat16_single_tile_size_bytes,
+        kNumOutputBeforeReductionTiles);
+
+    auto cb_output = create_circular_buffer(
+        program,
+        all_cores,
+        kOutputCbIndex,
+        data_format,
+        bfloat16_single_tile_size_bytes,
+        2 * block_size);  // create twice block_size variable
+
+    // -------------------------------------------------------------------------
+    // 3) Create reader/writer kernels
+    // -------------------------------------------------------------------------
+
+    auto* input_buffer = input.buffer();
+    TT_FATAL(
+        input_buffer->buffer_type() == ttnn::BufferType::DRAM,
+        "Input buffer must be in DRAM. Input buffer of type {}",
+        magic_enum::enum_name(input_buffer->buffer_type()));
+
+    auto* target_buffer = target.buffer();
+    TT_FATAL(
+        target_buffer->buffer_type() == ttnn::BufferType::DRAM,
+        "Target buffer must be in DRAM. Target buffer of type {}",
+        magic_enum::enum_name(target_buffer->buffer_type()));
+
+    auto* output_buffer = output.buffer();
+    TT_FATAL(
+        output_buffer->buffer_type() == ttnn::BufferType::DRAM,
+        "Output buffer must be in DRAM. Output buffer of type {}",
+        magic_enum::enum_name(output_buffer->buffer_type()));
+
+    // configure defines
+    std::map<std::string, std::string> defines;
+
+    // define whether I need mask or not
+    if (mask_w != 0) {
+        defines[kMaskWDefineKey] = "1";
+    }
+    // setup defines for reduce
+    // Compute kernel does not compile without these defines
+    // LLK reduction uses define values as default template parameters
+    defines["REDUCE_OP"] = "PoolType::SUM";
+    defines["REDUCE_DIM"] = "ReduceDim::REDUCE_ROW";
+
+    CrossEntropyForwardKernels kernels;
+    kernels.reader = create_reader_kernel(
+        program,
+        all_cores,
+        /* reader_compile_args */ {num_input_tiles, block_size, Wt, mask_w},
path: tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/cross_entropy_fw_program_factory.cpp
user: login: rfurko-tt
body: remove `num_input_tiles`
created_at: 2025-04-12T06:54:32Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,244 @@
+
+// SPDX-FileCopyrightText: (c) 2025 Tenstorrent AI ULC
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include <gtest/gtest.h>
+#include <sys/types.h>
+
+#include <cassert>
+#include <core/ttnn_all_includes.hpp>
+#include <cstddef>
+#include <ttnn/operations/reduction/generic/generic_reductions.hpp>
+
+#include "autograd/auto_context.hpp"
+#include "core/tt_tensor_utils.hpp"
+#include "metal/operations.hpp"
+
+class CrossEntropyForwardTest : public ::testing::Test {
+protected:
+    void SetUp() override {
+        ttml::autograd::ctx().open_device();
+    }
+
+    void TearDown() override {
+        ttml::autograd::ctx().close_device();
+    }
+};
+
+xt::xarray<float> calculate_cross_entropy_loss(const xt::xarray<float>& input, const xt::xarray<float>& target) {
+    xt::xarray<float> max_input = xt::amax(input, -1);
+    xt::xarray<float> max_input_expanded = xt::expand_dims(max_input, 3);
+    xt::xarray<float> shift_input = input - max_input_expanded;
+    xt::xarray<float> sum_exp_input = xt::sum(xt::exp(shift_input), -1);
+    xt::xarray<float> log_sum_exp_input = xt::log(sum_exp_input);
+    xt::xarray<float> log_sum_exp_input_expanded = xt::expand_dims(log_sum_exp_input, 3);
+    xt::xarray<float> result = -shift_input * target + target * log_sum_exp_input_expanded;
+    xt::xarray<float> reduced_result = xt::sum(result, -1, xt::keep_dims);
+    return reduced_result;
+}
+
+TEST_F(CrossEntropyForwardTest, CrossEntropyForward_Small_Forward) {
+    using namespace ttml;
+
+    const uint32_t N = 1, C = 1, H = 1, W = 8;
+
+    xt::xarray<float> example_xtensor = {{{{1.F, 2.F, 3.F, 4.F, 1.F, 2.F, 3.F, 4.F}}}};
+    auto input_logits = core::from_xtensor(example_xtensor, &autograd::ctx().get_device());
+    std::cout << "Input Logits:\n";
+    input_logits.print();
+
+    xt::xarray<float> target_xtensor = {{{{0.F, 1.F, 0.F, 0.F, 0.F, 0.F, 0.F, 0.F}}}};
+    assert((target_xtensor.shape() == example_xtensor.shape()));
+    auto target = core::from_xtensor(target_xtensor, &autograd::ctx().get_device());
+    std::cout << "Input Targets:\n";
+    target.print();
+
+    auto result = ttml::metal::cross_entropy_fw(input_logits, target);
+    std::cout << "CrossEntropyForward_Test:\nResult:\n";
+    result.print();
+
+    auto expected_result = calculate_cross_entropy_loss(example_xtensor, target_xtensor);
+    auto expected_result_print = core::from_xtensor(expected_result, &autograd::ctx().get_device());
+    std::cout << "Expected Result From :\n";
+    expected_result_print.print();
+
+    // Check if the result is close to the expected result
+    auto result_xtensor = core::to_xtensor(result);
+    assert((result_xtensor.shape() == expected_result.shape()));
+    EXPECT_TRUE(xt::allclose(result_xtensor, expected_result, 1e-2F));
+}
+
+TEST_F(CrossEntropyForwardTest, CrossEntropyForward_Negetive_Values) {
+    using namespace ttml;
+
+    const uint32_t N = 1, C = 1, H = 2, W = 4;
+
+    xt::xarray<float> example_xtensor = {{{{-100.F, -101.F, -102.F, -103.F}, {-5.01F, -5.02F, -0.3F, -7.F}}}};
+    auto input_logits = core::from_xtensor(example_xtensor, &autograd::ctx().get_device());
+    std::cout << "Input Logits:\n";
+    input_logits.print();
+
+    xt::xarray<float> target_xtensor = {{{{1.F, 0.F, 0.F, 0.F}, {0.F, 0.F, 1.F, 0.F}}}};
+    assert((target_xtensor.shape() == example_xtensor.shape()));
+    auto target = core::from_xtensor(target_xtensor, &autograd::ctx().get_device());
+    std::cout << "Input Targets:\n";
+    target.print();
+
+    auto result = ttml::metal::cross_entropy_fw(input_logits, target);
+    std::cout << "CrossEntropyForward_Test:\nResult:\n";
+    result.print();
+
+    auto expected_result = calculate_cross_entropy_loss(example_xtensor, target_xtensor);
+    auto expected_result_print = core::from_xtensor(expected_result, &autograd::ctx().get_device());
+    std::cout << "Expected Result From :\n";
+    expected_result_print.print();
+
+    // Check if the result is close to the expected result
+    auto result_xtensor = core::to_xtensor(result);
+    assert((result_xtensor.shape() == expected_result.shape()));
+    EXPECT_TRUE(xt::allclose(result_xtensor, expected_result, 1e-2F, 1e-2F));
+}
+
+TEST_F(CrossEntropyForwardTest, CrossEntropyForward_Batch) {
+    using namespace ttml;
+
+    const uint32_t N = 3, C = 1, H = 2, W = 47;
+    const auto shape = ttnn::SmallVector<uint32_t>{N, C, H, W};
+    std::vector<float> logits_data(N * C * H * W);
+    std::vector<float> target_data(N * C * H * W, 0.0f);
+
+    std::random_device rd;
+    std::mt19937 gen(rd());
+    std::uniform_real_distribution<float> dist(-10.0f, 10.0f);
+    std::uniform_int_distribution<std::size_t> class_dist(0, W - 1);
+    for (uint32_t n = 0; n < N; ++n) {
+        for (uint32_t h = 0; h < H; ++h) {
+            uint32_t offset = (n * H + h) * W;  // flattened offset
+            for (uint32_t w = 0; w < W; ++w) {
+                logits_data[offset + w] = dist(gen);  // random logits
+            }
+            std::size_t true_class = class_dist(gen) % W;
+            logits_data[offset + true_class] = 100 + n + h;  // one-hot
path: tt-train/tests/ops/cross_entropy_fw_op_test.cpp
user: login: rfurko-tt
body: not clear what this line is supposed to do
created_at: 2025-04-12T06:55:37Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,306 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include <compute_kernel_api/cb_api.h>
+#include <compute_kernel_api/pack.h>
+#include <compute_kernel_api/reg_api.h>
+#include <debug/dprint.h>
+#include <tensix.h>
+
+#include <cstdint>
+
+#include "compute_kernel_api.h"
+#include "compute_kernel_api/bcast.h"
+#include "compute_kernel_api/common.h"
+#include "compute_kernel_api/eltwise_binary.h"
+#include "compute_kernel_api/eltwise_binary_sfpu.h"
+#include "compute_kernel_api/eltwise_unary/eltwise_unary.h"
+#include "compute_kernel_api/eltwise_unary/exp.h"
+#include "compute_kernel_api/eltwise_unary/negative.h"
+#include "compute_kernel_api/eltwise_unary/recip.h"
+#include "compute_kernel_api/eltwise_unary/sfpu_split_includes.h"
+#include "compute_kernel_api/eltwise_unary/sqrt.h"
+#include "compute_kernel_api/mask.h"
+#include "compute_kernel_api/reduce.h"
+#include "compute_kernel_api/tile_move_copy.h"
+
+namespace NAMESPACE {
+
+constexpr uint32_t num_rows_per_core = get_compile_time_arg_val(0);  // rows to process in this kernel
+constexpr uint32_t block_size = get_compile_time_arg_val(1);         // size of block
+constexpr uint32_t Wt = get_compile_time_arg_val(2);                 // number of tiles in inner dimension
+
+constexpr auto cb_input = tt::CBIndex::c_0;
+constexpr auto cb_target = tt::CBIndex::c_1;
+constexpr auto cb_mask = tt::CBIndex::c_2;
+constexpr auto cb_max_mask = tt::CBIndex::c_3;
+constexpr auto cb_scaler = tt::CBIndex::c_4;  // used to reduction
+constexpr auto cb_max_value_before_reduction = tt::CBIndex::c_5;
+constexpr auto cb_max_value_after_reduction = tt::CBIndex::c_6;
+constexpr auto cb_exp_sum_before_reduction = tt::CBIndex::c_7;
+constexpr auto cb_exp_sum_after_reduction = tt::CBIndex::c_8;
+constexpr auto cb_output_before_reduction = tt::CBIndex::c_9;
+constexpr auto cb_output = tt::CBIndex::c_10;
+
+constexpr uint32_t onetile = 1;
+#ifdef DO_MASK_W
+constexpr bool do_mask_w = true;
+#else
+constexpr bool do_mask_w = false;
+#endif
+
+void find_max_value_in_row() {
+    const uint32_t max_value_register = 0;
+    const uint32_t tile_register = 1U;
+    cb_wait_front(cb_input, Wt);  // wait until reader kernel has written Wt tiles to input buffer
+    cb_reserve_back(cb_max_value_before_reduction, onetile);
+    tile_regs_acquire();
+    for (uint32_t col = 0; col < Wt; ++col) {
+        auto working_register = col == 0 ? max_value_register : tile_register;
+        copy_tile_init(cb_input);
+        copy_tile(cb_input, /* tile_idx */ col, /* register_idx */ working_register);
+
+        if constexpr (do_mask_w) {
+            if (col + 1 == Wt) {
+                // this is limitation of the function mask_tile
+                // mask tile currently does not work for mask register that is not next to data register
+                const uint32_t mask_register = working_register + 1U;  // mask register should be next to data register
+                copy_tile_init(cb_max_mask);
+                copy_tile(cb_max_mask, /* tile_idx */ 0, /* register idx */ mask_register);
+
+                add_binary_tile_init();
+                add_binary_tile(working_register, mask_register);
+            }
+        }
+
+        if (col > 0) {
+            max_tile_init();
+            max_tile(max_value_register, tile_register);
+        }
+    }
+    // cb_pop_front(cb_input, Wt);
+    tile_regs_commit();
+
+    tile_regs_wait();
+    pack_reconfig_data_format(cb_max_value_before_reduction);
+    pack_tile(max_value_register, cb_max_value_before_reduction);
+    tile_regs_release();
+    cb_push_back(cb_max_value_before_reduction, onetile);
+
+    // reduce max value and push it to cb_max_value_after_reduction
+    cb_wait_front(cb_max_value_before_reduction, onetile);
+    cb_reserve_back(cb_max_value_after_reduction, onetile);
+
+    const uint32_t reduction_register = 0;
+    tile_regs_acquire();
+    reconfig_data_format(cb_max_value_before_reduction, cb_scaler);
+    reduce_init_delta<false, PoolType::MAX, ReduceDim::REDUCE_ROW>(
+        cb_max_value_before_reduction, cb_scaler, cb_max_value_after_reduction);
+    reduce_tile<PoolType::MAX, ReduceDim::REDUCE_ROW>(
+        cb_max_value_before_reduction,
+        cb_scaler,
+        /* tile_idx */ 0,
+        /* tile_idx */ 0,
+        reduction_register);
+    reduce_revert_delta<ReduceDim::REDUCE_ROW>(cb_max_value_before_reduction);
+    tile_regs_commit();
+
+    tile_regs_wait();
+
+    pack_reconfig_data_format(cb_max_value_after_reduction);
+    pack_tile(reduction_register, cb_max_value_after_reduction);
+    tile_regs_release();
+
+    cb_push_back(cb_max_value_after_reduction, onetile);   // push tile after reduction
+    cb_pop_front(cb_max_value_before_reduction, onetile);  // pop tile before reduction
+}
+
+// calculate sum of exp(x - max(x))
+void calculate_sum_exp_x() {
+    // need to read from input second time
+    // cb_wait_front(cb_input, Wt);  // wait until reader kernel has written Wt tiles to input buffer
+    cb_wait_front(cb_max_value_after_reduction, onetile);  // wait until we get max value in each row
+
+    // run through all tiles in row
+    const uint32_t accum_register = 0;
+    const uint32_t tile_register = 1U;
+
+    cb_reserve_back(cb_exp_sum_before_reduction, onetile);  // create extra buffer
+    tile_regs_acquire();
+    for (uint32_t col = 0; col < Wt; ++col) {
+        auto working_register = col == 0 ? accum_register : tile_register;
+
+        // subtract max value from each tile
+        sub_bcast_cols_init_short(cb_input, cb_max_value_after_reduction);
+        sub_tiles_bcast_cols(
+            cb_input,
+            cb_max_value_after_reduction,
+            /* tile idx */ col,
+            /* tile idx */ 0,
+            /* reg tile idx */ working_register);
+
+        exp_tile_init();
+        exp_tile</* approx */ false>(working_register);  // calculate exp for each tile in tile register
+
+        if constexpr (do_mask_w) {
+            if (col + 1 == Wt) {
+                // this is limitation of the function mask_tile
+                // mask tile currently does not work for mask register that is not next to data register
+                const uint32_t mask_register = working_register + 1U;  // mask register should be next to data register
+                copy_tile_init(cb_mask);
+                copy_tile(cb_mask, /* tile_idx */ 0, /* register idx */ mask_register);
+
+                mask_tile_init();
+                mask_tile(working_register, mask_register);  // mask should be next to tile register
+            }
+        }
+
+        if (col > 0) {
+            add_binary_tile_init();
+            add_binary_tile(accum_register, working_register);
+        }
+    }
+    tile_regs_commit();
+
+    // packing to cb_exp_sum_before_reduction
+    tile_regs_wait();
+    pack_reconfig_data_format(cb_exp_sum_before_reduction);
+    pack_tile(accum_register, cb_exp_sum_before_reduction);
+    tile_regs_release();
+    cb_push_back(cb_exp_sum_before_reduction, onetile);
+
+    // cb_pop_front(cb_input, Wt);  // delete Wt tiles from input buffer
+
+    // exp(x - max(x))
+    // reduce exp sum to cb_exp_sum_after_reduction
+    cb_wait_front(cb_exp_sum_before_reduction, onetile);
+    cb_reserve_back(cb_exp_sum_after_reduction, onetile);
+
+    tile_regs_acquire();
+    const uint32_t reduction_register = 0;
+    reconfig_data_format(cb_exp_sum_before_reduction, cb_scaler);
+    reduce_init_delta<false, PoolType::SUM, ReduceDim::REDUCE_ROW>(
+        cb_exp_sum_before_reduction, cb_scaler, cb_exp_sum_after_reduction);
+    reduce_tile<PoolType::SUM, ReduceDim::REDUCE_ROW>(
+        cb_exp_sum_before_reduction,
+        cb_scaler,
+        /* tile_idx */ 0,
+        /* tile_idx */ 0,
+        /* reduction_register */ reduction_register);
+    reduce_revert_delta<ReduceDim::REDUCE_ROW>(cb_exp_sum_before_reduction);
+
+    log_tile_init();
+    log_tile(reduction_register);
+
+    tile_regs_commit();
+    tile_regs_wait();
+    pack_reconfig_data_format(cb_exp_sum_after_reduction);
+    pack_tile(/* tile_idx */ reduction_register, cb_exp_sum_after_reduction);
+    tile_regs_release();
+    cb_push_back(cb_exp_sum_after_reduction, onetile);
+    cb_pop_front(cb_exp_sum_before_reduction, onetile);
+}
+
+void MAIN {
+    if constexpr (do_mask_w) {
+        cb_wait_front(cb_mask, onetile);
+        cb_wait_front(cb_max_mask, onetile);
+    }
+    cb_wait_front(cb_scaler, onetile);
+
+    init_sfpu(cb_input, cb_output);
+    binary_op_init_common(cb_input, cb_target, cb_output);
+
+    for (uint32_t row = 0; row < num_rows_per_core; ++row) {
+        find_max_value_in_row();  // find max value in each row
+        calculate_sum_exp_x();    // calculate sum of exp(x - max(x))
+        //  set (input +  sum of exp(x - max())) to output
+        {
+            // need to read input third time
+            // cb_wait_front(cb_input, Wt);   // wait until reader kernel has written Wt tiles to input buffer
+            cb_wait_front(cb_target, Wt);  // wait until reader kernel has written Wt tiles to target buffer
+            cb_wait_front(cb_exp_sum_after_reduction, onetile);  // <- get here log(sum(exp(x - max(x)))
+
+            cb_reserve_back(
+                cb_output, onetile);  // reserve Wt tiles in output buffer == wait until cb will has Wt tiles
+            // for (uint32_t col = 0; col < Wt; col += block_size) {
+            const uint32_t accum_register = 0;
+            tile_regs_acquire();
+            reconfig_data_format(cb_input, cb_max_value_after_reduction);
+            // accumulate log(sum(exp(x - max(x))) over row
+            for (uint32_t col = 0; col < Wt; col++) {
+                auto working_register = col == 0 ? 0 : 1U;
+                auto target_register = working_register + 1U;
+                auto log_sum_exp_register = working_register + 2U;
+                // cb_reserve_back(cb_output, block_size);  // wait until output cb will have block_size free tiles
+                sub_bcast_cols_init_short(cb_input, cb_max_value_after_reduction);
+                sub_tiles_bcast_cols(
+                    cb_input,
+                    cb_max_value_after_reduction,
+                    /* tile idx */ col,
+                    /* tile idx */ 0,
+                    /* reg tile idx */ working_register);
+
+                mul_bcast_cols_init_short(cb_target, cb_exp_sum_after_reduction);
+                mul_tiles_bcast_cols(
path: tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/kernels/compute/cross_entropy_fw_kernel.cpp
user: login: rfurko-tt
body: we can multiply by target once (in this code it's twice)
created_at: 2025-04-12T07:01:13Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,244 @@
+
path: tt-train/tests/ops/cross_entropy_fw_op_test.cpp
user: login: rfurko-tt
body: remove empty line
created_at: 2025-04-12T07:03:24Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,17 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include "cross_entropy_fw.hpp"
+
+#include "device/cross_entropy_fw_device_operation.hpp"
+
+namespace ttml::metal::ops::cross_entropy_fw {
+
+ttnn::Tensor CrossEntropyForwardOperation::invoke(const ttnn::Tensor& input_tensor, const ttnn::Tensor& target_tensor) {
+    auto result /*ttnn::Tensor*/ = ttnn::prim::ttml_cross_entropy_fw(
path: tt-train/sources/ttml/metal/ops/cross_entropy_fw/cross_entropy_fw.cpp
user: login: vmelnykovTT
body: 
Yeah, the function return value also indicates the type of the result
created_at: 2025-04-16T14:13:09Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,474 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include "cross_entropy_fw_program_factory.hpp"
+
+#include <bit>
+#include <cstdint>
+
+#include "cross_entropy_fw_device_operation_types.hpp"
+
+namespace {
+
+constexpr auto kReaderKernelPath =
+    "tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/kernels/dataflow/"
+    "reader_cross_entropy_fw_interleaved_start_id.cpp";
+
+constexpr auto kWriterKernelPath =
+    "tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/kernels/dataflow/"
+    "writer_cross_entropy_fw_interleaved_start_id.cpp";
+
+constexpr auto kComputeKernelPath =
+    "tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/kernels/compute/cross_entropy_fw_kernel.cpp";
+
+// reader runtime args
+constexpr uint32_t kInputBufferIdx = 0;
+constexpr uint32_t kTargetBufferIdx = 1U;
+// writer runtime args
+constexpr uint32_t kOutputBufferIdx = 0;
+
+constexpr auto kInputCbIndex = tt::CBIndex::c_0;
+constexpr auto kTargetCbIndex = tt::CBIndex::c_1;
+constexpr auto kMaskCbIndex = tt::CBIndex::c_2;
+constexpr auto kMaxMaskCbIndex = tt::CBIndex::c_3;
+constexpr auto kScalerCbIndex = tt::CBIndex::c_4;  // used to reduction
+constexpr auto kMaxValueBeforeReductionCbIndex = tt::CBIndex::c_5;
+constexpr auto kMaxValueAfterReductionCbIndex = tt::CBIndex::c_6;
+constexpr auto kExpSumBeforeReductionCbIndex = tt::CBIndex::c_7;
+constexpr auto KExpSumAfterReductionCbIndex = tt::CBIndex::c_8;
+constexpr auto kOutputBeforeReductionCbIndex = tt::CBIndex::c_9;
+constexpr auto kOutputCbIndex = tt::CBIndex::c_10;
+
+constexpr uint32_t kNumMaskTiles = 1U;
+constexpr uint32_t kNumMaxValueTiles = 1U;
+constexpr uint32_t kNumOutputTiles = 1U;
+constexpr uint32_t kMaxValueBeforeReductionTiles = 1U;
+constexpr uint32_t kNumMaxValueAfterReductionTiles = 1U;
+constexpr uint32_t kNumExpSumBeforeReductionTiles = 1U;
+constexpr uint32_t kNumExpSumAfterReductionTiles = 1U;
+constexpr uint32_t kNumScalerTiles = 1U;  // used it to reduction ???
+constexpr uint32_t kNumOutputBeforeReductionTiles = 1U;
+
+const std::string kMaskWDefineKey = "DO_MASK_W";
+
+uint32_t get_block_size(uint32_t num_inner) {
+    const uint32_t max_block_size = 4U;  // 4 is the maximum block size for enabled fp32 dest acc
+    for (uint32_t block_size = max_block_size; block_size > 1; block_size--) {
+        if (num_inner % block_size == 0) {  // if num_inner is divisible by block_size - choose this block_size
+            return block_size;
+        }
+    }
+    return 1U;
+}
+
+}  // namespace
+
+namespace ttml::metal::ops::cross_entropy_fw::device {
+
+/**
+ *   Helper struct to hold references to all kernels we create,
+ *        used during runtime argument setup.
+ */
+struct CrossEntropyForwardKernels {
+    tt::tt_metal::KernelHandle reader;
+    tt::tt_metal::KernelHandle writer;
+    tt::tt_metal::KernelHandle compute_group_1;
+    tt::tt_metal::KernelHandle compute_group_2;
+};
+
+/**
+ *   Create and configure a circular buffer, returning both the configuration and the handle.
+ */
+tt::tt_metal::CBHandle create_circular_buffer(
+    tt::tt_metal::Program& program,
+    const tt::tt_metal::CoreRangeSet& core_ranges,
+    uint32_t cb_index,
+    tt::DataFormat data_format,
+    uint32_t single_tile_size,
+    uint32_t num_tiles) {
+    tt::tt_metal::CircularBufferConfig cb_config =
+        tt::tt_metal::CircularBufferConfig(num_tiles * single_tile_size, {{cb_index, data_format}})
+            .set_page_size(cb_index, single_tile_size);
+
+    auto cb_handle = CreateCircularBuffer(program, core_ranges, cb_config);
+    return cb_handle;
+}
+
+/**
+ *   Create a reader kernel with the given compile-time arguments.
+ */
+tt::tt_metal::KernelHandle create_reader_kernel(
+    tt::tt_metal::Program& program,
+    const tt::tt_metal::CoreRangeSet& core_ranges,
+    const std::vector<uint32_t>& compile_time_args,
+    const std::map<std::string, std::string>& defines,
+    const std::string& kernel_path) {
+    return tt::tt_metal::CreateKernel(
+        program, kernel_path, core_ranges, tt::tt_metal::ReaderDataMovementConfig(compile_time_args, defines));
+}
+
+/**
+ *   Create a writer kernel with the given compile-time arguments.
+ */
+tt::tt_metal::KernelHandle create_writer_kernel(
+    tt::tt_metal::Program& program,
+    const tt::tt_metal::CoreRangeSet& core_ranges,
+    const std::vector<uint32_t>& compile_time_args,
+    const std::map<std::string, std::string>& defines,
+    const std::string& kernel_path) {
+    return tt::tt_metal::CreateKernel(
+        program, kernel_path, core_ranges, tt::tt_metal::WriterDataMovementConfig(compile_time_args, defines));
+}
+
+/**
+ * Create a compute kernel with the given compile-time arguments.
+ */
+tt::tt_metal::KernelHandle create_compute_kernel(
+    tt::tt_metal::Program& program,
+    const tt::tt_metal::CoreRangeSet& core_ranges,
+    const std::vector<uint32_t>& compile_time_args,
+    const std::map<std::string, std::string>& defines,
+    const std::string& kernel_path) {
+    return tt::tt_metal::CreateKernel(
+        program,
+        kernel_path,
+        core_ranges,
+        tt::tt_metal::ComputeConfig{
+            .math_fidelity = MathFidelity::HiFi4,
+            .fp32_dest_acc_en = true,
+            .math_approx_mode = false,
+            .compile_args = compile_time_args,
+            .defines = defines});
+}
+
+/**
+ * Set up the runtime arguments for the 4 relevant kernels (reader, writer, compute G1, compute G2)
+ *        for each core in the grid.
+ */
+void assign_per_core_runtime_args(
+    tt::tt_metal::Program& program,
+    const CrossEntropyForwardKernels& kernels,
+    const tt::tt_metal::Buffer* input_buffer,
+    const tt::tt_metal::Buffer* target_buffer,
+    const tt::tt_metal::Buffer* output_buffer,
+    uint32_t num_cores,
+    uint32_t num_cores_y,
+    uint32_t num_rows_per_core_group_1,
+    uint32_t num_rows_per_core_group_2,
+    const tt::tt_metal::CoreRangeSet& core_group_1,
+    const tt::tt_metal::CoreRangeSet& core_group_2) {
+    for (uint32_t i = 0, num_rows_written = 0; i < num_cores; i++) {
+        CoreCoord core = {i / num_cores_y, i % num_cores_y};
+
+        // Determine how many rows this core will process
+        uint32_t num_rows_per_core = 0;
+        if (core_group_1.contains(core)) {
+            num_rows_per_core = num_rows_per_core_group_1;
+        } else if (core_group_2.contains(core)) {
+            num_rows_per_core = num_rows_per_core_group_2;
+        } else {
+            TT_FATAL(false, "Core not in specified core ranges");
+        }
+
+        // Reader kernel: (input_addr, gamma_addr, number_of_rows, offset_in_rows)
+        SetRuntimeArgs(
+            program,
+            kernels.reader,
+            core,
+            {input_buffer->address(), target_buffer->address(), num_rows_per_core, num_rows_written});
+
+        // Writer kernel: (dst_addr, dst_rms_addr number_of_rows, offset_in_rows)
+        SetRuntimeArgs(
+            program,
+            kernels.writer,
+            core,
+            {output_buffer->address(), /*temporary*/ num_rows_per_core, num_rows_written});
+
+        num_rows_written += num_rows_per_core;
+    }
+}
+
+CrossEntropyForwardProgramFactory::cached_program_t CrossEntropyForwardProgramFactory::create(
+    const operation_attributes_t& operation_attributes,
+    const tensor_args_t& tensor_args,
+    tensor_return_value_t& output) {
+    // -------------------------------------------------------------------------
+    // 1) Setup device, data formats, tile sizes, and compute split
+    // -------------------------------------------------------------------------
+    const auto& input = tensor_args.input;
+    const auto& target = tensor_args.target;
+
+    auto* device = input.device();
+
+    tt::tt_metal::Program program{};
+
+    tt::DataFormat input_data_format = datatype_to_dataformat_converter(input.get_dtype());
+    TT_FATAL(input_data_format == tt::DataFormat::Float16_b, "Input data format must be Float16_b");
+
+    uint32_t bfloat16_single_tile_size_bytes = tt::tt_metal::detail::TileSize(tt::DataFormat::Float16_b);
+    uint32_t float32_single_tile_size_bytes = tt::tt_metal::detail::TileSize(tt::DataFormat::Float32);
+
+    auto padded_tensor_shape = input.get_padded_shape();
+    auto padded_tensor_volume = input.volume();
+
+    TT_FATAL(
+        padded_tensor_volume % tt::constants::TILE_HW == 0, "Padded input tensor volume must be divisible by TILE_HW");
+    TT_FATAL(padded_tensor_shape.rank() == 4U, "Input tensor must be 4D");
+
+    //
+    uint32_t Wt = padded_tensor_shape[-1] / tt::constants::TILE_WIDTH;  // <- number of tiles in inner dimension
+    uint32_t Ht = padded_tensor_shape[-2] / tt::constants::TILE_HEIGHT;
+    uint32_t NC = padded_tensor_shape[0] * padded_tensor_shape[1];
+    uint32_t total_rows_to_process = NC * Ht;
+
+    // get number of free cores
+    auto compute_with_storage_grid_size = device->compute_with_storage_grid_size();
+    uint32_t num_cores_x = compute_with_storage_grid_size.x;
+    uint32_t num_cores_y = compute_with_storage_grid_size.y;
+
+    // get the number of inner dimension(D) - we will caclulate log softmax on this dimension
+    uint32_t num_inner = input.get_logical_shape()[-1];  // (B, 1, S, D) and we take D here
+
+    // mask_w - this mask used to avoid calculation of extra data(data which will be added to create full tile 32x32)??
+    uint32_t mask_w = num_inner % tt::constants::TILE_WIDTH;  // <- width index of first trash value in tile
+
+    // compile arguments
+    uint32_t block_size = get_block_size(Wt);
+
+    auto [num_cores, all_cores, core_group_1, core_group_2, num_rows_per_core_group_1, num_rows_per_core_group_2] =
+        tt::tt_metal::split_work_to_cores(compute_with_storage_grid_size, total_rows_to_process);
+
+    // -------------------------------------------------------------------------
+    // 2) Create and configure circular buffers
+    // -------------------------------------------------------------------------
+
+    // TODO: add calculation of required  memory and check if it fits in L1
+    // then add possibility to kernels to process data by blocks
+    const uint32_t available_L1_in_bytes =
+        device->l1_size_per_core() - device->allocator()->get_base_allocator_addr(tt::tt_metal::HalMemType::L1);
+
+    const uint64_t mask_scaler_maxmask_memory =
+        (kNumMaskTiles + kNumScalerTiles + kNumMaskTiles) * bfloat16_single_tile_size_bytes;
+
+    const uint32_t num_input_tiles = Wt;
+    const uint32_t num_target_tiles = Wt;
+
+    auto data_format = input_data_format;  // tt::DataFormat::Float16_b
+    auto precise_data_format = tt::DataFormat::Float32;
+
+    auto cb_input = create_circular_buffer(
+        program, all_cores, kInputCbIndex, data_format, bfloat16_single_tile_size_bytes, num_input_tiles);
+
+    auto cb_target = create_circular_buffer(
+        program, all_cores, kTargetCbIndex, data_format, bfloat16_single_tile_size_bytes, num_target_tiles);
+
+    auto cb_mask = create_circular_buffer(
+        program, all_cores, kMaskCbIndex, data_format, bfloat16_single_tile_size_bytes, kNumMaskTiles);
+
+    auto cb_max_mask = create_circular_buffer(
+        program, all_cores, kMaxMaskCbIndex, data_format, bfloat16_single_tile_size_bytes, kNumMaskTiles);
+
+    auto cb_scaler = create_circular_buffer(
+        program, all_cores, kScalerCbIndex, data_format, bfloat16_single_tile_size_bytes, kNumScalerTiles);
+
+    auto cb_max_value_before_reduction = create_circular_buffer(
+        program,
+        all_cores,
+        kMaxValueBeforeReductionCbIndex,
+        data_format,
+        bfloat16_single_tile_size_bytes,
+        kMaxValueBeforeReductionTiles);
+
+    auto cb_max_value_after_reduction = create_circular_buffer(
+        program,
+        all_cores,
+        kMaxValueAfterReductionCbIndex,
+        data_format,
+        bfloat16_single_tile_size_bytes,
+        kNumMaxValueAfterReductionTiles);
+
+    auto cb_exp_sum_before_reduction = create_circular_buffer(
+        program,
+        all_cores,
+        kExpSumBeforeReductionCbIndex,
+        data_format,
+        bfloat16_single_tile_size_bytes,
+        kNumExpSumBeforeReductionTiles);
+
+    auto cb_exp_sum_after_refuction = create_circular_buffer(
+        program,
+        all_cores,
+        KExpSumAfterReductionCbIndex,
+        data_format,
+        bfloat16_single_tile_size_bytes,
+        kNumExpSumAfterReductionTiles);
+
+    auto cb_output_before_refuction = create_circular_buffer(
+        program,
+        all_cores,
+        kOutputBeforeReductionCbIndex,
+        data_format,
+        bfloat16_single_tile_size_bytes,
+        kNumOutputBeforeReductionTiles);
+
+    auto cb_output = create_circular_buffer(
+        program,
+        all_cores,
+        kOutputCbIndex,
+        data_format,
+        bfloat16_single_tile_size_bytes,
+        2 * block_size);  // create twice block_size variable
path: tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/cross_entropy_fw_program_factory.cpp
user: login: vmelnykovTT
body: I only need one tile here, since I get the result and write it to the output once
created_at: 2025-04-16T14:25:35Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,306 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include <compute_kernel_api/cb_api.h>
+#include <compute_kernel_api/pack.h>
+#include <compute_kernel_api/reg_api.h>
+#include <debug/dprint.h>
+#include <tensix.h>
+
+#include <cstdint>
+
+#include "compute_kernel_api.h"
+#include "compute_kernel_api/bcast.h"
+#include "compute_kernel_api/common.h"
+#include "compute_kernel_api/eltwise_binary.h"
+#include "compute_kernel_api/eltwise_binary_sfpu.h"
+#include "compute_kernel_api/eltwise_unary/eltwise_unary.h"
+#include "compute_kernel_api/eltwise_unary/exp.h"
+#include "compute_kernel_api/eltwise_unary/negative.h"
+#include "compute_kernel_api/eltwise_unary/recip.h"
+#include "compute_kernel_api/eltwise_unary/sfpu_split_includes.h"
+#include "compute_kernel_api/eltwise_unary/sqrt.h"
+#include "compute_kernel_api/mask.h"
+#include "compute_kernel_api/reduce.h"
+#include "compute_kernel_api/tile_move_copy.h"
+
+namespace NAMESPACE {
+
+constexpr uint32_t num_rows_per_core = get_compile_time_arg_val(0);  // rows to process in this kernel
+constexpr uint32_t block_size = get_compile_time_arg_val(1);         // size of block
+constexpr uint32_t Wt = get_compile_time_arg_val(2);                 // number of tiles in inner dimension
+
+constexpr auto cb_input = tt::CBIndex::c_0;
+constexpr auto cb_target = tt::CBIndex::c_1;
+constexpr auto cb_mask = tt::CBIndex::c_2;
+constexpr auto cb_max_mask = tt::CBIndex::c_3;
+constexpr auto cb_scaler = tt::CBIndex::c_4;  // used to reduction
+constexpr auto cb_max_value_before_reduction = tt::CBIndex::c_5;
+constexpr auto cb_max_value_after_reduction = tt::CBIndex::c_6;
+constexpr auto cb_exp_sum_before_reduction = tt::CBIndex::c_7;
+constexpr auto cb_exp_sum_after_reduction = tt::CBIndex::c_8;
+constexpr auto cb_output_before_reduction = tt::CBIndex::c_9;
+constexpr auto cb_output = tt::CBIndex::c_10;
+
+constexpr uint32_t onetile = 1;
+#ifdef DO_MASK_W
+constexpr bool do_mask_w = true;
+#else
+constexpr bool do_mask_w = false;
+#endif
+
+void find_max_value_in_row() {
+    const uint32_t max_value_register = 0;
+    const uint32_t tile_register = 1U;
+    cb_wait_front(cb_input, Wt);  // wait until reader kernel has written Wt tiles to input buffer
+    cb_reserve_back(cb_max_value_before_reduction, onetile);
+    tile_regs_acquire();
+    for (uint32_t col = 0; col < Wt; ++col) {
+        auto working_register = col == 0 ? max_value_register : tile_register;
+        copy_tile_init(cb_input);
+        copy_tile(cb_input, /* tile_idx */ col, /* register_idx */ working_register);
+
+        if constexpr (do_mask_w) {
+            if (col + 1 == Wt) {
+                // this is limitation of the function mask_tile
+                // mask tile currently does not work for mask register that is not next to data register
+                const uint32_t mask_register = working_register + 1U;  // mask register should be next to data register
+                copy_tile_init(cb_max_mask);
+                copy_tile(cb_max_mask, /* tile_idx */ 0, /* register idx */ mask_register);
+
+                add_binary_tile_init();
+                add_binary_tile(working_register, mask_register);
+            }
+        }
+
+        if (col > 0) {
+            max_tile_init();
+            max_tile(max_value_register, tile_register);
+        }
+    }
+    // cb_pop_front(cb_input, Wt);
+    tile_regs_commit();
+
+    tile_regs_wait();
+    pack_reconfig_data_format(cb_max_value_before_reduction);
+    pack_tile(max_value_register, cb_max_value_before_reduction);
+    tile_regs_release();
+    cb_push_back(cb_max_value_before_reduction, onetile);
+
+    // reduce max value and push it to cb_max_value_after_reduction
+    cb_wait_front(cb_max_value_before_reduction, onetile);
+    cb_reserve_back(cb_max_value_after_reduction, onetile);
+
+    const uint32_t reduction_register = 0;
+    tile_regs_acquire();
+    reconfig_data_format(cb_max_value_before_reduction, cb_scaler);
+    reduce_init_delta<false, PoolType::MAX, ReduceDim::REDUCE_ROW>(
+        cb_max_value_before_reduction, cb_scaler, cb_max_value_after_reduction);
+    reduce_tile<PoolType::MAX, ReduceDim::REDUCE_ROW>(
+        cb_max_value_before_reduction,
+        cb_scaler,
+        /* tile_idx */ 0,
+        /* tile_idx */ 0,
+        reduction_register);
+    reduce_revert_delta<ReduceDim::REDUCE_ROW>(cb_max_value_before_reduction);
+    tile_regs_commit();
+
+    tile_regs_wait();
+
+    pack_reconfig_data_format(cb_max_value_after_reduction);
+    pack_tile(reduction_register, cb_max_value_after_reduction);
+    tile_regs_release();
+
+    cb_push_back(cb_max_value_after_reduction, onetile);   // push tile after reduction
+    cb_pop_front(cb_max_value_before_reduction, onetile);  // pop tile before reduction
+}
+
+// calculate sum of exp(x - max(x))
+void calculate_sum_exp_x() {
+    // need to read from input second time
+    // cb_wait_front(cb_input, Wt);  // wait until reader kernel has written Wt tiles to input buffer
+    cb_wait_front(cb_max_value_after_reduction, onetile);  // wait until we get max value in each row
+
+    // run through all tiles in row
+    const uint32_t accum_register = 0;
+    const uint32_t tile_register = 1U;
+
+    cb_reserve_back(cb_exp_sum_before_reduction, onetile);  // create extra buffer
+    tile_regs_acquire();
+    for (uint32_t col = 0; col < Wt; ++col) {
+        auto working_register = col == 0 ? accum_register : tile_register;
+
+        // subtract max value from each tile
+        sub_bcast_cols_init_short(cb_input, cb_max_value_after_reduction);
+        sub_tiles_bcast_cols(
+            cb_input,
+            cb_max_value_after_reduction,
+            /* tile idx */ col,
+            /* tile idx */ 0,
+            /* reg tile idx */ working_register);
+
+        exp_tile_init();
+        exp_tile</* approx */ false>(working_register);  // calculate exp for each tile in tile register
+
+        if constexpr (do_mask_w) {
+            if (col + 1 == Wt) {
+                // this is limitation of the function mask_tile
+                // mask tile currently does not work for mask register that is not next to data register
+                const uint32_t mask_register = working_register + 1U;  // mask register should be next to data register
+                copy_tile_init(cb_mask);
+                copy_tile(cb_mask, /* tile_idx */ 0, /* register idx */ mask_register);
+
+                mask_tile_init();
+                mask_tile(working_register, mask_register);  // mask should be next to tile register
+            }
+        }
+
+        if (col > 0) {
+            add_binary_tile_init();
+            add_binary_tile(accum_register, working_register);
+        }
+    }
+    tile_regs_commit();
+
+    // packing to cb_exp_sum_before_reduction
+    tile_regs_wait();
+    pack_reconfig_data_format(cb_exp_sum_before_reduction);
+    pack_tile(accum_register, cb_exp_sum_before_reduction);
+    tile_regs_release();
+    cb_push_back(cb_exp_sum_before_reduction, onetile);
+
+    // cb_pop_front(cb_input, Wt);  // delete Wt tiles from input buffer
+
+    // exp(x - max(x))
+    // reduce exp sum to cb_exp_sum_after_reduction
+    cb_wait_front(cb_exp_sum_before_reduction, onetile);
+    cb_reserve_back(cb_exp_sum_after_reduction, onetile);
+
+    tile_regs_acquire();
+    const uint32_t reduction_register = 0;
+    reconfig_data_format(cb_exp_sum_before_reduction, cb_scaler);
+    reduce_init_delta<false, PoolType::SUM, ReduceDim::REDUCE_ROW>(
+        cb_exp_sum_before_reduction, cb_scaler, cb_exp_sum_after_reduction);
+    reduce_tile<PoolType::SUM, ReduceDim::REDUCE_ROW>(
+        cb_exp_sum_before_reduction,
+        cb_scaler,
+        /* tile_idx */ 0,
+        /* tile_idx */ 0,
+        /* reduction_register */ reduction_register);
+    reduce_revert_delta<ReduceDim::REDUCE_ROW>(cb_exp_sum_before_reduction);
+
+    log_tile_init();
+    log_tile(reduction_register);
+
+    tile_regs_commit();
+    tile_regs_wait();
+    pack_reconfig_data_format(cb_exp_sum_after_reduction);
+    pack_tile(/* tile_idx */ reduction_register, cb_exp_sum_after_reduction);
+    tile_regs_release();
+    cb_push_back(cb_exp_sum_after_reduction, onetile);
+    cb_pop_front(cb_exp_sum_before_reduction, onetile);
+}
+
+void MAIN {
+    if constexpr (do_mask_w) {
+        cb_wait_front(cb_mask, onetile);
+        cb_wait_front(cb_max_mask, onetile);
+    }
+    cb_wait_front(cb_scaler, onetile);
+
+    init_sfpu(cb_input, cb_output);
+    binary_op_init_common(cb_input, cb_target, cb_output);
+
+    for (uint32_t row = 0; row < num_rows_per_core; ++row) {
+        find_max_value_in_row();  // find max value in each row
+        calculate_sum_exp_x();    // calculate sum of exp(x - max(x))
+        //  set (input +  sum of exp(x - max())) to output
+        {
+            // need to read input third time
+            // cb_wait_front(cb_input, Wt);   // wait until reader kernel has written Wt tiles to input buffer
+            cb_wait_front(cb_target, Wt);  // wait until reader kernel has written Wt tiles to target buffer
+            cb_wait_front(cb_exp_sum_after_reduction, onetile);  // <- get here log(sum(exp(x - max(x)))
+
+            cb_reserve_back(
+                cb_output, onetile);  // reserve Wt tiles in output buffer == wait until cb will has Wt tiles
+            // for (uint32_t col = 0; col < Wt; col += block_size) {
+            const uint32_t accum_register = 0;
+            tile_regs_acquire();
+            reconfig_data_format(cb_input, cb_max_value_after_reduction);
+            // accumulate log(sum(exp(x - max(x))) over row
+            for (uint32_t col = 0; col < Wt; col++) {
+                auto working_register = col == 0 ? 0 : 1U;
+                auto target_register = working_register + 1U;
+                auto log_sum_exp_register = working_register + 2U;
+                // cb_reserve_back(cb_output, block_size);  // wait until output cb will have block_size free tiles
+                sub_bcast_cols_init_short(cb_input, cb_max_value_after_reduction);
+                sub_tiles_bcast_cols(
+                    cb_input,
+                    cb_max_value_after_reduction,
+                    /* tile idx */ col,
+                    /* tile idx */ 0,
+                    /* reg tile idx */ working_register);
+
+                mul_bcast_cols_init_short(cb_target, cb_exp_sum_after_reduction);
+                mul_tiles_bcast_cols(
path: tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/kernels/compute/cross_entropy_fw_kernel.cpp
user: login: vmelnykovTT
body: Changed the calculation approach and completely avoided broadcasting
created_at: 2025-04-16T14:26:40Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,17 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#pragma once
+
+#include "metal/ttnn_all_includes.hpp"
+
+namespace ttml::metal::ops::cross_entropy_fw {
+
+struct CrossEntropyForwardOperation {
+    static ttnn::Tensor invoke(
+        const ttnn::Tensor& input_tensor,  // logits : model output (N, 1, H, W)
path: tt-train/sources/ttml/metal/ops/cross_entropy_fw/cross_entropy_fw.hpp
user: login: rfurko-tt
body: what do you think about removing `_tensor` as it is clearly has type ttnn::Tensor
created_at: 2025-04-24T18:34:46Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,17 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#pragma once
+
+#include "metal/ttnn_all_includes.hpp"
+
+namespace ttml::metal::ops::cross_entropy_fw {
+
+struct CrossEntropyForwardOperation {
+    static ttnn::Tensor invoke(
+        const ttnn::Tensor& input_tensor,  // logits : model output (N, 1, H, W)
path: tt-train/sources/ttml/metal/ops/cross_entropy_fw/cross_entropy_fw.hpp
user: login: rfurko-tt
body: Same for tensor below
created_at: 2025-04-24T18:34:52Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,133 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include "cross_entropy_fw_device_operation.hpp"
+
+#include "cross_entropy_fw_program_factory.hpp"
+
+namespace ttml::metal::ops::cross_entropy_fw::device {
+
+CrossEntropyForwardDeviceOperation::program_factory_t CrossEntropyForwardDeviceOperation::select_program_factory(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    return CrossEntropyForwardProgramFactory{};
+}
+
+void CrossEntropyForwardDeviceOperation::validate_on_program_cache_hit(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    validate_on_program_cache_miss(args, tensor_args);
+}
+
+void CrossEntropyForwardDeviceOperation::validate_on_program_cache_miss(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    auto check_tensor = [](const ttnn::Tensor& tensor,
+                           const std::string& name,
+                           const tt::tt_metal::Layout required_layout,
+                           const tt::tt_metal::DataType required_dtype) {
+        TT_FATAL(
+            tensor.device()->arch() == tt::ARCH::WORMHOLE_B0,
+            "CrossEntropyForward operation is only supported on Wormhole. Device arch: {}. Tensor name {}",
+            magic_enum::enum_name(tensor.device()->arch()),
+            name);
+
+        TT_FATAL(
+            tensor.storage_type() == tt::tt_metal::StorageType::DEVICE,
+            "CrossEntropyForward operation requires {} to be on Device. Input storage type: {}",
+            name,
+            static_cast<int>(tensor.storage_type()));
+
+        TT_FATAL(
+            tensor.buffer() != nullptr,
+            "Operands to CrossEntropyForward need to be allocated in buffers on the device. Buffer is null. Tensor "
+            "name {}",
+            name);
+
+        TT_FATAL(
+            tensor.get_layout() == required_layout,
+            "CrossEntropyForward operation requires tensor to be in Tile layout. {} tensor layout: {}",
+            name,
+            static_cast<int>(tensor.get_layout()));
+
+        TT_FATAL(
+            tensor.get_dtype() == required_dtype,
+            "CrossEntropyForward operation requires tensor to be of BFLOAT16 data type. {} tensor data type: {}",
+            name,
+            static_cast<int>(tensor.get_dtype()));
+
+        TT_FATAL(
+            tensor.memory_config().memory_layout == ttnn::TensorMemoryLayout::INTERLEAVED,
+            "CrossEntropyForward operation requires Interleaved memory layout. {} "
+            "memory layout: `{}`",
+            name,
+            static_cast<int>(tensor.memory_config().memory_layout));
path: tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/cross_entropy_fw_device_operation.cpp
user: login: rfurko-tt
body: could you please use `magic_enum::enum_name` for error message?
created_at: 2025-04-24T18:35:21Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,133 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include "cross_entropy_fw_device_operation.hpp"
+
+#include "cross_entropy_fw_program_factory.hpp"
+
+namespace ttml::metal::ops::cross_entropy_fw::device {
+
+CrossEntropyForwardDeviceOperation::program_factory_t CrossEntropyForwardDeviceOperation::select_program_factory(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    return CrossEntropyForwardProgramFactory{};
+}
+
+void CrossEntropyForwardDeviceOperation::validate_on_program_cache_hit(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    validate_on_program_cache_miss(args, tensor_args);
+}
+
+void CrossEntropyForwardDeviceOperation::validate_on_program_cache_miss(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    auto check_tensor = [](const ttnn::Tensor& tensor,
+                           const std::string& name,
+                           const tt::tt_metal::Layout required_layout,
+                           const tt::tt_metal::DataType required_dtype) {
+        TT_FATAL(
+            tensor.device()->arch() == tt::ARCH::WORMHOLE_B0,
+            "CrossEntropyForward operation is only supported on Wormhole. Device arch: {}. Tensor name {}",
+            magic_enum::enum_name(tensor.device()->arch()),
+            name);
+
+        TT_FATAL(
+            tensor.storage_type() == tt::tt_metal::StorageType::DEVICE,
+            "CrossEntropyForward operation requires {} to be on Device. Input storage type: {}",
+            name,
+            static_cast<int>(tensor.storage_type()));
+
+        TT_FATAL(
+            tensor.buffer() != nullptr,
+            "Operands to CrossEntropyForward need to be allocated in buffers on the device. Buffer is null. Tensor "
+            "name {}",
+            name);
+
+        TT_FATAL(
+            tensor.get_layout() == required_layout,
+            "CrossEntropyForward operation requires tensor to be in Tile layout. {} tensor layout: {}",
+            name,
+            static_cast<int>(tensor.get_layout()));
+
+        TT_FATAL(
+            tensor.get_dtype() == required_dtype,
+            "CrossEntropyForward operation requires tensor to be of BFLOAT16 data type. {} tensor data type: {}",
+            name,
+            static_cast<int>(tensor.get_dtype()));
+
+        TT_FATAL(
+            tensor.memory_config().memory_layout == ttnn::TensorMemoryLayout::INTERLEAVED,
+            "CrossEntropyForward operation requires Interleaved memory layout. {} "
+            "memory layout: `{}`",
+            name,
+            static_cast<int>(tensor.memory_config().memory_layout));
+    };
+
+    const auto& input_tensor = tensor_args.input;
+    const auto& target_tensor = tensor_args.target;
+    const auto& preallocated_output_tensor = tensor_args.preallocated_output;
+    check_tensor(input_tensor, "Input", tt::tt_metal::Layout::TILE, tt::tt_metal::DataType::BFLOAT16);
+    check_tensor(target_tensor, "Target", tt::tt_metal::Layout::ROW_MAJOR, tt::tt_metal::DataType::UINT32);
+    if (preallocated_output_tensor.has_value()) {
+        check_tensor(
+            preallocated_output_tensor.value(),
+            "Preallocated Output",
+            tt::tt_metal::Layout::TILE,
+            tt::tt_metal::DataType::BFLOAT16);
+    }
+}
+
+CrossEntropyForwardDeviceOperation::spec_return_value_t CrossEntropyForwardDeviceOperation::compute_output_specs(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    if (tensor_args.preallocated_output.has_value()) {
+        return tensor_args.preallocated_output->get_tensor_spec();
+    }
+    auto input_logical_shape = tensor_args.input.get_logical_shape();
+    input_logical_shape[-1] = 1;
path: tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/cross_entropy_fw_device_operation.cpp
user: login: rfurko-tt
body: = 1U;
created_at: 2025-04-24T18:35:31Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,133 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include "cross_entropy_fw_device_operation.hpp"
+
+#include "cross_entropy_fw_program_factory.hpp"
+
+namespace ttml::metal::ops::cross_entropy_fw::device {
+
+CrossEntropyForwardDeviceOperation::program_factory_t CrossEntropyForwardDeviceOperation::select_program_factory(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    return CrossEntropyForwardProgramFactory{};
+}
+
+void CrossEntropyForwardDeviceOperation::validate_on_program_cache_hit(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    validate_on_program_cache_miss(args, tensor_args);
+}
+
+void CrossEntropyForwardDeviceOperation::validate_on_program_cache_miss(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    auto check_tensor = [](const ttnn::Tensor& tensor,
+                           const std::string& name,
+                           const tt::tt_metal::Layout required_layout,
+                           const tt::tt_metal::DataType required_dtype) {
+        TT_FATAL(
+            tensor.device()->arch() == tt::ARCH::WORMHOLE_B0,
+            "CrossEntropyForward operation is only supported on Wormhole. Device arch: {}. Tensor name {}",
+            magic_enum::enum_name(tensor.device()->arch()),
+            name);
+
+        TT_FATAL(
+            tensor.storage_type() == tt::tt_metal::StorageType::DEVICE,
+            "CrossEntropyForward operation requires {} to be on Device. Input storage type: {}",
+            name,
+            static_cast<int>(tensor.storage_type()));
+
+        TT_FATAL(
+            tensor.buffer() != nullptr,
+            "Operands to CrossEntropyForward need to be allocated in buffers on the device. Buffer is null. Tensor "
+            "name {}",
+            name);
+
+        TT_FATAL(
+            tensor.get_layout() == required_layout,
+            "CrossEntropyForward operation requires tensor to be in Tile layout. {} tensor layout: {}",
+            name,
+            static_cast<int>(tensor.get_layout()));
+
+        TT_FATAL(
+            tensor.get_dtype() == required_dtype,
+            "CrossEntropyForward operation requires tensor to be of BFLOAT16 data type. {} tensor data type: {}",
+            name,
+            static_cast<int>(tensor.get_dtype()));
+
+        TT_FATAL(
+            tensor.memory_config().memory_layout == ttnn::TensorMemoryLayout::INTERLEAVED,
+            "CrossEntropyForward operation requires Interleaved memory layout. {} "
+            "memory layout: `{}`",
+            name,
+            static_cast<int>(tensor.memory_config().memory_layout));
+    };
+
+    const auto& input_tensor = tensor_args.input;
+    const auto& target_tensor = tensor_args.target;
+    const auto& preallocated_output_tensor = tensor_args.preallocated_output;
+    check_tensor(input_tensor, "Input", tt::tt_metal::Layout::TILE, tt::tt_metal::DataType::BFLOAT16);
+    check_tensor(target_tensor, "Target", tt::tt_metal::Layout::ROW_MAJOR, tt::tt_metal::DataType::UINT32);
+    if (preallocated_output_tensor.has_value()) {
+        check_tensor(
+            preallocated_output_tensor.value(),
+            "Preallocated Output",
+            tt::tt_metal::Layout::TILE,
+            tt::tt_metal::DataType::BFLOAT16);
+    }
+}
+
+CrossEntropyForwardDeviceOperation::spec_return_value_t CrossEntropyForwardDeviceOperation::compute_output_specs(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    if (tensor_args.preallocated_output.has_value()) {
+        return tensor_args.preallocated_output->get_tensor_spec();
+    }
+    auto input_logical_shape = tensor_args.input.get_logical_shape();
+    input_logical_shape[-1] = 1;
+    return ttnn::TensorSpec(
+        ttnn::Shape(input_logical_shape),
+        tt::tt_metal::TensorLayout(
+            tensor_args.input.get_dtype(), tt::tt_metal::Layout::TILE, tensor_args.input.memory_config()));
+}
+
+CrossEntropyForwardDeviceOperation::tensor_return_value_t CrossEntropyForwardDeviceOperation::create_output_tensors(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    tensor_return_value_t output_tensor;
+
+    spec_return_value_t output_specs = compute_output_specs(args, tensor_args);
+
+    if (tensor_args.preallocated_output.has_value()) {
+        output_tensor = tensor_args.preallocated_output.value();
+    } else {
+        output_tensor = create_device_tensor(output_specs, tensor_args.input.device());
+    }
+
+    return output_tensor;
+}
+
+tt::stl::hash::hash_t CrossEntropyForwardDeviceOperation::compute_program_hash(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    const auto& input_tensor = tensor_args.input;
+    const auto& input_logical_shape = input_tensor.get_logical_shape();
+    auto program_factory = select_program_factory(args, tensor_args);
+    tt::tt_metal::operation::Hash hash = tt::tt_metal::operation::hash_operation<CrossEntropyForwardDeviceOperation>(
path: tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/cross_entropy_fw_device_operation.cpp
user: login: rfurko-tt
body: auto hash = ...;
created_at: 2025-04-24T18:37:29Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,133 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include "cross_entropy_fw_device_operation.hpp"
+
+#include "cross_entropy_fw_program_factory.hpp"
+
+namespace ttml::metal::ops::cross_entropy_fw::device {
+
+CrossEntropyForwardDeviceOperation::program_factory_t CrossEntropyForwardDeviceOperation::select_program_factory(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    return CrossEntropyForwardProgramFactory{};
+}
+
+void CrossEntropyForwardDeviceOperation::validate_on_program_cache_hit(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    validate_on_program_cache_miss(args, tensor_args);
+}
+
+void CrossEntropyForwardDeviceOperation::validate_on_program_cache_miss(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    auto check_tensor = [](const ttnn::Tensor& tensor,
+                           const std::string& name,
+                           const tt::tt_metal::Layout required_layout,
+                           const tt::tt_metal::DataType required_dtype) {
+        TT_FATAL(
+            tensor.device()->arch() == tt::ARCH::WORMHOLE_B0,
+            "CrossEntropyForward operation is only supported on Wormhole. Device arch: {}. Tensor name {}",
+            magic_enum::enum_name(tensor.device()->arch()),
+            name);
+
+        TT_FATAL(
+            tensor.storage_type() == tt::tt_metal::StorageType::DEVICE,
+            "CrossEntropyForward operation requires {} to be on Device. Input storage type: {}",
+            name,
+            static_cast<int>(tensor.storage_type()));
+
+        TT_FATAL(
+            tensor.buffer() != nullptr,
+            "Operands to CrossEntropyForward need to be allocated in buffers on the device. Buffer is null. Tensor "
+            "name {}",
+            name);
+
+        TT_FATAL(
+            tensor.get_layout() == required_layout,
+            "CrossEntropyForward operation requires tensor to be in Tile layout. {} tensor layout: {}",
+            name,
+            static_cast<int>(tensor.get_layout()));
+
+        TT_FATAL(
+            tensor.get_dtype() == required_dtype,
+            "CrossEntropyForward operation requires tensor to be of BFLOAT16 data type. {} tensor data type: {}",
+            name,
+            static_cast<int>(tensor.get_dtype()));
+
+        TT_FATAL(
+            tensor.memory_config().memory_layout == ttnn::TensorMemoryLayout::INTERLEAVED,
+            "CrossEntropyForward operation requires Interleaved memory layout. {} "
+            "memory layout: `{}`",
+            name,
+            static_cast<int>(tensor.memory_config().memory_layout));
+    };
+
+    const auto& input_tensor = tensor_args.input;
+    const auto& target_tensor = tensor_args.target;
+    const auto& preallocated_output_tensor = tensor_args.preallocated_output;
+    check_tensor(input_tensor, "Input", tt::tt_metal::Layout::TILE, tt::tt_metal::DataType::BFLOAT16);
+    check_tensor(target_tensor, "Target", tt::tt_metal::Layout::ROW_MAJOR, tt::tt_metal::DataType::UINT32);
+    if (preallocated_output_tensor.has_value()) {
+        check_tensor(
+            preallocated_output_tensor.value(),
+            "Preallocated Output",
+            tt::tt_metal::Layout::TILE,
+            tt::tt_metal::DataType::BFLOAT16);
+    }
+}
+
+CrossEntropyForwardDeviceOperation::spec_return_value_t CrossEntropyForwardDeviceOperation::compute_output_specs(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    if (tensor_args.preallocated_output.has_value()) {
+        return tensor_args.preallocated_output->get_tensor_spec();
+    }
+    auto input_logical_shape = tensor_args.input.get_logical_shape();
+    input_logical_shape[-1] = 1;
+    return ttnn::TensorSpec(
+        ttnn::Shape(input_logical_shape),
+        tt::tt_metal::TensorLayout(
+            tensor_args.input.get_dtype(), tt::tt_metal::Layout::TILE, tensor_args.input.memory_config()));
+}
+
+CrossEntropyForwardDeviceOperation::tensor_return_value_t CrossEntropyForwardDeviceOperation::create_output_tensors(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    tensor_return_value_t output_tensor;
+
+    spec_return_value_t output_specs = compute_output_specs(args, tensor_args);
+
+    if (tensor_args.preallocated_output.has_value()) {
+        output_tensor = tensor_args.preallocated_output.value();
+    } else {
+        output_tensor = create_device_tensor(output_specs, tensor_args.input.device());
+    }
+
+    return output_tensor;
+}
+
+tt::stl::hash::hash_t CrossEntropyForwardDeviceOperation::compute_program_hash(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    const auto& input_tensor = tensor_args.input;
+    const auto& input_logical_shape = input_tensor.get_logical_shape();
+    auto program_factory = select_program_factory(args, tensor_args);
+    tt::tt_metal::operation::Hash hash = tt::tt_metal::operation::hash_operation<CrossEntropyForwardDeviceOperation>(
+        args, program_factory.index(), input_tensor.dtype(), input_logical_shape);
+
+    return hash;
+}
+
+std::
+    tuple<CrossEntropyForwardDeviceOperation::operation_attributes_t, CrossEntropyForwardDeviceOperation::tensor_args_t>
path: tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/cross_entropy_fw_device_operation.cpp
user: login: rfurko-tt
body: move to a single line, don't like this way
created_at: 2025-04-24T18:37:50Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,508 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include "cross_entropy_fw_program_factory.hpp"
+
+#include <bit>
+#include <cstdint>
+#include <tt-metalium/buffer.hpp>
+
+#include "cross_entropy_fw_device_operation_types.hpp"
+
+namespace {
+
+constexpr auto kReaderKernelPath =
+    "tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/kernels/dataflow/"
+    "reader_cross_entropy_fw_interleaved_start_id.cpp";
+
+constexpr auto kWriterKernelPath =
+    "tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/kernels/dataflow/"
+    "writer_cross_entropy_fw_interleaved_start_id.cpp";
+
+constexpr auto kComputeKernelPath =
+    "tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/kernels/compute/cross_entropy_fw_kernel.cpp";
+
+// reader runtime args
+constexpr uint32_t kInputBufferIdx = 0;
+constexpr uint32_t kTargetBufferIdx = 1U;
+// writer runtime args
+constexpr uint32_t kOutputBufferIdx = 0;
+
+constexpr auto kInputCbIndex = tt::CBIndex::c_0;
+constexpr auto kTargetCbIndex = tt::CBIndex::c_1;
+constexpr auto kMaskCbIndex = tt::CBIndex::c_2;
+constexpr auto kMaxMaskCbIndex = tt::CBIndex::c_3;
+constexpr auto kScalerCbIndex = tt::CBIndex::c_4;  // used to reduction
+constexpr auto kInputTileCbIndex = tt::CBIndex::c_5;
+constexpr auto kTargetLogitsCbIndex = tt::CBIndex::c_6;
+constexpr auto kMaxValueBeforeReductionCbIndex = tt::CBIndex::c_7;
+constexpr auto kMaxValueAfterReductionCbIndex = tt::CBIndex::c_8;
+constexpr auto kExpSumBeforeReductionCbIndex = tt::CBIndex::c_9;
+constexpr auto KExpSumAfterReductionCbIndex = tt::CBIndex::c_10;
+constexpr auto kOutputCbIndex = tt::CBIndex::c_11;
+
+constexpr uint32_t kNumTargetIndexesTiles = 1U;
+constexpr uint32_t kNumMaskTiles = 1U;
+constexpr uint32_t kNumMaxValueTiles = 1U;
+constexpr uint32_t kNumInputTilesByIndx = 1U;
+constexpr uint32_t kMaxValueBeforeReductionTiles = 1U;
+constexpr uint32_t kNumTargetLogitsTiles = 1U;
+constexpr uint32_t kNumMaxValueAfterReductionTiles = 1U;
path: tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/cross_entropy_fw_program_factory.cpp
user: login: rfurko-tt
body: change a lot of this NumTiles to 2U, so you can prepare one additional while calculating 
created_at: 2025-04-24T18:39:03Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,508 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include "cross_entropy_fw_program_factory.hpp"
+
+#include <bit>
+#include <cstdint>
+#include <tt-metalium/buffer.hpp>
+
+#include "cross_entropy_fw_device_operation_types.hpp"
+
+namespace {
+
+constexpr auto kReaderKernelPath =
+    "tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/kernels/dataflow/"
+    "reader_cross_entropy_fw_interleaved_start_id.cpp";
+
+constexpr auto kWriterKernelPath =
+    "tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/kernels/dataflow/"
+    "writer_cross_entropy_fw_interleaved_start_id.cpp";
+
+constexpr auto kComputeKernelPath =
+    "tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/kernels/compute/cross_entropy_fw_kernel.cpp";
+
+// reader runtime args
+constexpr uint32_t kInputBufferIdx = 0;
+constexpr uint32_t kTargetBufferIdx = 1U;
+// writer runtime args
+constexpr uint32_t kOutputBufferIdx = 0;
+
+constexpr auto kInputCbIndex = tt::CBIndex::c_0;
+constexpr auto kTargetCbIndex = tt::CBIndex::c_1;
+constexpr auto kMaskCbIndex = tt::CBIndex::c_2;
+constexpr auto kMaxMaskCbIndex = tt::CBIndex::c_3;
+constexpr auto kScalerCbIndex = tt::CBIndex::c_4;  // used to reduction
+constexpr auto kInputTileCbIndex = tt::CBIndex::c_5;
+constexpr auto kTargetLogitsCbIndex = tt::CBIndex::c_6;
+constexpr auto kMaxValueBeforeReductionCbIndex = tt::CBIndex::c_7;
+constexpr auto kMaxValueAfterReductionCbIndex = tt::CBIndex::c_8;
+constexpr auto kExpSumBeforeReductionCbIndex = tt::CBIndex::c_9;
+constexpr auto KExpSumAfterReductionCbIndex = tt::CBIndex::c_10;
+constexpr auto kOutputCbIndex = tt::CBIndex::c_11;
+
+constexpr uint32_t kNumTargetIndexesTiles = 1U;
+constexpr uint32_t kNumMaskTiles = 1U;
+constexpr uint32_t kNumMaxValueTiles = 1U;
+constexpr uint32_t kNumInputTilesByIndx = 1U;
+constexpr uint32_t kMaxValueBeforeReductionTiles = 1U;
+constexpr uint32_t kNumTargetLogitsTiles = 1U;
+constexpr uint32_t kNumMaxValueAfterReductionTiles = 1U;
+constexpr uint32_t kNumExpSumBeforeReductionTiles = 1U;
+constexpr uint32_t kNumExpSumAfterReductionTiles = 1U;
+constexpr uint32_t kNumScalerTiles = 1U;  // used it to reduction ???
+constexpr uint32_t kNumOutputBeforeReductionTiles = 1U;
+constexpr uint32_t kNumOutputTiles = 1U;
+
+constexpr uint32_t kPageElementsNumber = 32U;
+
+const std::string kMaskWDefineKey = "DO_MASK_W";
+const std::string kEverythingFitsInL1DefineKey = "EVERYTHING_FITS_IN_L1";
+
+uint32_t get_block_size(uint32_t num_inner) {
+    const uint32_t max_block_size = 4U;  // 4 is the maximum block size for enabled fp32 dest acc
+    for (uint32_t block_size = max_block_size; block_size > 1; block_size--) {
+        if (num_inner % block_size == 0) {  // if num_inner is divisible by block_size - choose this block_size
+            return block_size;
+        }
+    }
+    return 1U;
+}
+
+}  // namespace
+
+namespace ttml::metal::ops::cross_entropy_fw::device {
+
+/**
+ *   Helper struct to hold references to all kernels we create,
+ *        used during runtime argument setup.
+ */
+struct CrossEntropyForwardKernels {
+    tt::tt_metal::KernelHandle reader;
+    tt::tt_metal::KernelHandle writer;
+    tt::tt_metal::KernelHandle compute_group_1;
+    tt::tt_metal::KernelHandle compute_group_2;
+};
+
+/**
+ *   Create and configure a circular buffer, returning both the configuration and the handle.
+ */
+tt::tt_metal::CBHandle create_circular_buffer(
+    tt::tt_metal::Program& program,
+    const tt::tt_metal::CoreRangeSet& core_ranges,
+    uint32_t cb_index,
+    tt::DataFormat data_format,
+    uint32_t single_tile_size,
+    uint32_t num_tiles) {
+    tt::tt_metal::CircularBufferConfig cb_config =
+        tt::tt_metal::CircularBufferConfig(num_tiles * single_tile_size, {{cb_index, data_format}})
+            .set_page_size(cb_index, single_tile_size);
+
+    auto cb_handle = CreateCircularBuffer(program, core_ranges, cb_config);
+    return cb_handle;
+}
+
+/**
+ *   Create a reader kernel with the given compile-time arguments.
+ */
+tt::tt_metal::KernelHandle create_reader_kernel(
+    tt::tt_metal::Program& program,
+    const tt::tt_metal::CoreRangeSet& core_ranges,
+    const std::vector<uint32_t>& compile_time_args,
+    const std::map<std::string, std::string>& defines,
+    const std::string& kernel_path) {
+    return tt::tt_metal::CreateKernel(
+        program, kernel_path, core_ranges, tt::tt_metal::ReaderDataMovementConfig(compile_time_args, defines));
+}
+
+/**
+ *   Create a writer kernel with the given compile-time arguments.
+ */
+tt::tt_metal::KernelHandle create_writer_kernel(
+    tt::tt_metal::Program& program,
+    const tt::tt_metal::CoreRangeSet& core_ranges,
+    const std::vector<uint32_t>& compile_time_args,
+    const std::map<std::string, std::string>& defines,
+    const std::string& kernel_path) {
+    return tt::tt_metal::CreateKernel(
+        program, kernel_path, core_ranges, tt::tt_metal::WriterDataMovementConfig(compile_time_args, defines));
+}
+
+/**
+ * Create a compute kernel with the given compile-time arguments.
+ */
+tt::tt_metal::KernelHandle create_compute_kernel(
+    tt::tt_metal::Program& program,
+    const tt::tt_metal::CoreRangeSet& core_ranges,
+    const std::vector<uint32_t>& compile_time_args,
+    const std::map<std::string, std::string>& defines,
+    const std::string& kernel_path) {
+    return tt::tt_metal::CreateKernel(
+        program,
+        kernel_path,
+        core_ranges,
+        tt::tt_metal::ComputeConfig{
+            .math_fidelity = MathFidelity::HiFi4,
+            .fp32_dest_acc_en = true,
+            .math_approx_mode = false,
+            .compile_args = compile_time_args,
+            .defines = defines});
+}
+
+/**
+ * Set up the runtime arguments for the 4 relevant kernels (reader, writer, compute G1, compute G2)
+ *        for each core in the grid.
+ */
+void assign_per_core_runtime_args(
+    tt::tt_metal::Program& program,
+    const CrossEntropyForwardKernels& kernels,
+    const tt::tt_metal::Buffer* input_buffer,
+    const tt::tt_metal::Buffer* target_buffer,
+    const tt::tt_metal::Buffer* output_buffer,
+
+    uint32_t num_cores,
+    uint32_t num_cores_y,
+    uint32_t num_rows_per_core_group_1,
+    uint32_t num_rows_per_core_group_2,
+    const tt::tt_metal::CoreRangeSet& core_group_1,
+    const tt::tt_metal::CoreRangeSet& core_group_2) {
+    for (uint32_t i = 0, num_rows_written = 0; i < num_cores; i++) {
+        CoreCoord core = {i / num_cores_y, i % num_cores_y};
+
+        // Determine how many rows this core will process
+        uint32_t num_rows_per_core = 0;
+        if (core_group_1.contains(core)) {
+            num_rows_per_core = num_rows_per_core_group_1;
+        } else if (core_group_2.contains(core)) {
+            num_rows_per_core = num_rows_per_core_group_2;
+        } else {
+            TT_FATAL(false, "Core not in specified core ranges");
+        }
+
+        // Reader kernel: (input_addr, gamma_addr, number_of_rows, offset_in_rows)
+        SetRuntimeArgs(
+            program,
+            kernels.reader,
+            core,
+            {input_buffer->address(), target_buffer->address(), num_rows_per_core, num_rows_written});
+
+        // Writer kernel: (dst_addr, dst_rms_addr number_of_rows, offset_in_rows)
+        SetRuntimeArgs(
+            program,
+            kernels.writer,
+            core,
+            {output_buffer->address(), /*temporary*/ num_rows_per_core, num_rows_written});
+
+        num_rows_written += num_rows_per_core;
+    }
+}
+
+CrossEntropyForwardProgramFactory::cached_program_t CrossEntropyForwardProgramFactory::create(
+    const operation_attributes_t& operation_attributes,
+    const tensor_args_t& tensor_args,
+    tensor_return_value_t& output) {
+    // -------------------------------------------------------------------------
+    // 1) Setup device, data formats, tile sizes, and compute split
+    // -------------------------------------------------------------------------
+    const auto& input = tensor_args.input;
+    const auto& target = tensor_args.target;
+
+    auto* device = input.device();
+
+    tt::tt_metal::Program program{};
+
+    tt::DataFormat input_data_format = datatype_to_dataformat_converter(input.get_dtype());
+    TT_FATAL(input_data_format == tt::DataFormat::Float16_b, "Input data format must be Float16_b");
+
+    uint32_t bfloat16_single_tile_size_bytes = tt::tt_metal::detail::TileSize(tt::DataFormat::Float16_b);
+    uint32_t float32_single_tile_size_bytes = tt::tt_metal::detail::TileSize(tt::DataFormat::Float32);
+
+    auto padded_tensor_shape = input.get_padded_shape();
+    auto padded_tensor_volume = input.volume();
+
+    TT_FATAL(
+        padded_tensor_volume % tt::constants::TILE_HW == 0, "Padded input tensor volume must be divisible by TILE_HW");
+    TT_FATAL(padded_tensor_shape.rank() == 4U, "Input tensor must be 4D");
+
+    //
+    uint32_t Wt = padded_tensor_shape[-1] / tt::constants::TILE_WIDTH;  // <- number of tiles in inner dimension
+    uint32_t Ht = padded_tensor_shape[-2] / tt::constants::TILE_HEIGHT;
+    uint32_t NC = padded_tensor_shape[0] * padded_tensor_shape[1];
+    uint32_t total_rows_to_process = NC * Ht;
+
+    // get size of target indexes inner dimension
+    uint32_t target_indexes_inner_dim_size = target.get_logical_shape()[-1] * target.element_size();
+    // read target indexes by pages(32 indexes in page)
+    uint32_t uint32_read_page_size = tt::datum_size(tt::DataFormat::UInt32) * kPageElementsNumber;
+    // tiled height of target indexes
+    uint32_t tiled_H =
path: tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/cross_entropy_fw_program_factory.cpp
user: login: rfurko-tt
body: i think you can use Ht 
created_at: 2025-04-24T18:40:00Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,404 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include <compute_kernel_api/cb_api.h>
+#include <compute_kernel_api/pack.h>
+#include <compute_kernel_api/reconfig_data_format.h>
+#include <compute_kernel_api/reg_api.h>
+#include <debug/dprint.h>
+#include <hostdevcommon/kernel_structs.h>
+#include <tensix.h>
+
+#include <cstdint>
+
+#include "compute_kernel_api.h"
+#include "compute_kernel_api/bcast.h"
+#include "compute_kernel_api/common.h"
+#include "compute_kernel_api/eltwise_binary.h"
+#include "compute_kernel_api/eltwise_binary_sfpu.h"
+#include "compute_kernel_api/eltwise_unary/eltwise_unary.h"
+#include "compute_kernel_api/eltwise_unary/exp.h"
+#include "compute_kernel_api/eltwise_unary/negative.h"
+#include "compute_kernel_api/eltwise_unary/recip.h"
+#include "compute_kernel_api/eltwise_unary/sfpu_split_includes.h"
+#include "compute_kernel_api/eltwise_unary/sqrt.h"
+#include "compute_kernel_api/mask.h"
+#include "compute_kernel_api/reduce.h"
+#include "compute_kernel_api/tile_move_copy.h"
+
+namespace NAMESPACE {
+
+constexpr uint32_t num_rows_per_core = get_compile_time_arg_val(0);  // rows to process in this kernel
+constexpr uint32_t block_size = get_compile_time_arg_val(1);         // size of block
+constexpr uint32_t Wt = get_compile_time_arg_val(2);                 // number of tiles in inner dimension
+
+constexpr auto cb_input = tt::CBIndex::c_0;
+constexpr auto cb_mask = tt::CBIndex::c_2;
+constexpr auto cb_max_mask = tt::CBIndex::c_3;
+constexpr auto cb_scaler = tt::CBIndex::c_4;  // used to reduction
+constexpr auto cb_target_logits = tt::CBIndex::c_6;
+constexpr auto cb_max_value_before_reduction = tt::CBIndex::c_7;
+constexpr auto cb_max_value_after_reduction = tt::CBIndex::c_8;
+constexpr auto cb_exp_sum_before_reduction = tt::CBIndex::c_9;
+constexpr auto cb_exp_sum_after_reduction = tt::CBIndex::c_10;
+constexpr auto cb_output = tt::CBIndex::c_11;
+
+constexpr uint32_t onetile = 1;
+
+#ifdef DO_MASK_W
+constexpr bool do_mask_w = true;
+#else
+constexpr bool do_mask_w = false;
+#endif
+
+#ifdef EVERYTHING_FITS_IN_L1
+
+void find_max_value_in_row() {
+    const uint32_t max_value_register = 0;
+    const uint32_t tile_register = 1U;
+    cb_wait_front(cb_input, Wt);  // wait until reader kernel has written Wt tiles to input buffer
+    cb_reserve_back(cb_max_value_before_reduction, onetile);
+    tile_regs_acquire();
+    for (uint32_t col = 0; col < Wt; ++col) {
+        auto working_register = col == 0 ? max_value_register : tile_register;
+        copy_tile_init(cb_input);
+        copy_tile(cb_input, /* tile_idx */ col, /* register_idx */ working_register);
+
+        if constexpr (do_mask_w) {
+            if (col + 1 == Wt) {
+                // this is limitation of the function mask_tile
+                // mask tile currently does not work for mask register that is not next to data register
+                const uint32_t mask_register = working_register + 1U;  // mask register should be next to data register
+
+                // NEXT 4 LINES ARE IMPORTANT AS WE FLASH WHAT IS INSIDE TRASH PADDING. IT CAN HAPPEN THAT VALUE INSIDE
path: tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/kernels/compute/cross_entropy_fw_kernel.cpp
user: login: rfurko-tt
body: improve this comment and do not use UPPER CASE :) 
created_at: 2025-04-24T18:41:35Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,404 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include <compute_kernel_api/cb_api.h>
+#include <compute_kernel_api/pack.h>
+#include <compute_kernel_api/reconfig_data_format.h>
+#include <compute_kernel_api/reg_api.h>
+#include <debug/dprint.h>
+#include <hostdevcommon/kernel_structs.h>
+#include <tensix.h>
+
+#include <cstdint>
+
+#include "compute_kernel_api.h"
+#include "compute_kernel_api/bcast.h"
+#include "compute_kernel_api/common.h"
+#include "compute_kernel_api/eltwise_binary.h"
+#include "compute_kernel_api/eltwise_binary_sfpu.h"
+#include "compute_kernel_api/eltwise_unary/eltwise_unary.h"
+#include "compute_kernel_api/eltwise_unary/exp.h"
+#include "compute_kernel_api/eltwise_unary/negative.h"
+#include "compute_kernel_api/eltwise_unary/recip.h"
+#include "compute_kernel_api/eltwise_unary/sfpu_split_includes.h"
+#include "compute_kernel_api/eltwise_unary/sqrt.h"
+#include "compute_kernel_api/mask.h"
+#include "compute_kernel_api/reduce.h"
+#include "compute_kernel_api/tile_move_copy.h"
+
+namespace NAMESPACE {
+
+constexpr uint32_t num_rows_per_core = get_compile_time_arg_val(0);  // rows to process in this kernel
+constexpr uint32_t block_size = get_compile_time_arg_val(1);         // size of block
+constexpr uint32_t Wt = get_compile_time_arg_val(2);                 // number of tiles in inner dimension
+
+constexpr auto cb_input = tt::CBIndex::c_0;
+constexpr auto cb_mask = tt::CBIndex::c_2;
+constexpr auto cb_max_mask = tt::CBIndex::c_3;
+constexpr auto cb_scaler = tt::CBIndex::c_4;  // used to reduction
+constexpr auto cb_target_logits = tt::CBIndex::c_6;
+constexpr auto cb_max_value_before_reduction = tt::CBIndex::c_7;
+constexpr auto cb_max_value_after_reduction = tt::CBIndex::c_8;
+constexpr auto cb_exp_sum_before_reduction = tt::CBIndex::c_9;
+constexpr auto cb_exp_sum_after_reduction = tt::CBIndex::c_10;
+constexpr auto cb_output = tt::CBIndex::c_11;
+
+constexpr uint32_t onetile = 1;
+
+#ifdef DO_MASK_W
+constexpr bool do_mask_w = true;
+#else
+constexpr bool do_mask_w = false;
+#endif
+
+#ifdef EVERYTHING_FITS_IN_L1
+
+void find_max_value_in_row() {
+    const uint32_t max_value_register = 0;
+    const uint32_t tile_register = 1U;
+    cb_wait_front(cb_input, Wt);  // wait until reader kernel has written Wt tiles to input buffer
+    cb_reserve_back(cb_max_value_before_reduction, onetile);
+    tile_regs_acquire();
+    for (uint32_t col = 0; col < Wt; ++col) {
+        auto working_register = col == 0 ? max_value_register : tile_register;
+        copy_tile_init(cb_input);
+        copy_tile(cb_input, /* tile_idx */ col, /* register_idx */ working_register);
+
+        if constexpr (do_mask_w) {
+            if (col + 1 == Wt) {
+                // this is limitation of the function mask_tile
+                // mask tile currently does not work for mask register that is not next to data register
+                const uint32_t mask_register = working_register + 1U;  // mask register should be next to data register
+
+                // NEXT 4 LINES ARE IMPORTANT AS WE FLASH WHAT IS INSIDE TRASH PADDING. IT CAN HAPPEN THAT VALUE INSIDE
+                // IS NAN, SO NAN+(-INF)=NAN AND NOT -INF AS WE EXPECTED. AS WELL AS -INF * 0 = NAN
+                copy_tile_init(cb_mask);
+                copy_tile(cb_mask, /* tile_idx */ 0, /* register idx */ mask_register);
+                mask_tile_init();
+                mask_tile(working_register, mask_register);  // mask should be next to tile register.
+
+                copy_tile_init(cb_max_mask);
+                copy_tile(cb_max_mask, /* tile_idx */ 0, /* register idx */ mask_register);
+
+                add_binary_tile_init();
+                add_binary_tile(working_register, mask_register);
+            }
+        }
+
+        if (col > 0) {
+            max_tile_init();
+            max_tile(max_value_register, tile_register);
+        }
+    }
+    tile_regs_commit();
+
+    tile_regs_wait();
+    pack_reconfig_data_format(cb_max_value_before_reduction);
+    pack_tile(max_value_register, cb_max_value_before_reduction);
+    tile_regs_release();
+    cb_push_back(cb_max_value_before_reduction, onetile);
+}
+
+#else
+
+void find_max_value_in_row() {
+    cb_reserve_back(cb_max_value_before_reduction, onetile);
+    tile_regs_acquire();
+    const uint32_t max_value_register = 0;
+    const uint32_t tile_register = 1U;
+    for (uint32_t col = 0; col < Wt;) {
+        cb_wait_front(cb_input, block_size);  // wait until reader kernel has written block_size tiles to input buffer
+        for (uint32_t block_idx = 0; block_idx < block_size; ++block_idx, ++col) {
+            auto working_register = col == 0 ? max_value_register : tile_register;
+            copy_tile_init(cb_input);
+            copy_tile(cb_input, /* tile_idx */ block_idx, /* register_idx */ working_register);
+
+            if constexpr (do_mask_w) {
+                if (col + 1 == Wt) {
+                    // this is limitation of the function mask_tile
+                    // mask tile currently does not work for mask register that is not next to data register
+                    const uint32_t mask_register =
+                        working_register + 1U;  // mask register should be next to data register
+
+                    // NEXT 4 LINES ARE IMPORTANT AS WE FLASH WHAT IS INSIDE TRASH PADDING. IT CAN HAPPEN THAT VALUE
path: tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/kernels/compute/cross_entropy_fw_kernel.cpp
user: login: rfurko-tt
body: Same here

created_at: 2025-04-24T18:42:31Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,227 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include <dataflow_api_addrgen.h>
+
+#include <cstdint>
+#include <cstring>
+
+#include "dataflow_api.h"
+#include "debug/dprint.h"
+#include "debug/dprint_pages.h"
+
+constexpr uint32_t FACE_HEIGHT = 16;
+constexpr uint32_t FACE_WIDTH = 16;
+constexpr uint32_t TILE_HEIGHT = 32;
+constexpr uint32_t TILE_WIDTH = 32;
+constexpr uint32_t TARGET_INDEXES_ELEMENTS_NUMBER = 32;
+
+float bfloat16_to_float(uint16_t bf16_bits) {
+    uint32_t full_bits = static_cast<uint32_t>(bf16_bits) << 16;
+    float result;
+    std::memcpy(&result, &full_bits, sizeof(result));
+    return result;
+}
+
+// calculate page and offset for target indexes
+std::pair<uint32_t, uint32_t> get_page_and_offset(uint32_t tiled_row, uint32_t tiled_H) {
+    // tasks number = N * C * (H / 32)
+    // tiled_row = start_row + i
path: tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/kernels/dataflow/reader_cross_entropy_fw_interleaved_start_id.cpp
user: login: rfurko-tt
body: remove this comments
created_at: 2025-04-24T18:45:52Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,37 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include "dataflow_api.h"
+
+void kernel_main() {
+    uint32_t runtime_args_counter = 0;
+    uint32_t output_addr = get_arg_val<uint32_t>(runtime_args_counter++);
+    uint32_t num_rows_to_process = get_arg_val<uint32_t>(runtime_args_counter++);
+    uint32_t start_row = get_arg_val<uint32_t>(runtime_args_counter++);
+
+    constexpr uint32_t cb_output_idx = tt::CBIndex::c_11;
+
+    constexpr uint32_t block_size = get_compile_time_arg_val(0);
+    constexpr uint32_t Wt = get_compile_time_arg_val(1);  // number of tiles in inner dimension
+
+    constexpr uint32_t onetile = 1U;
+
+    const uint32_t tile_bytes = get_tile_size(cb_output_idx);
+    const DataFormat data_format = get_dataformat(cb_output_idx);
+
+    const InterleavedAddrGenFast</* is dram */ true> output_addr_generator = {
+        .bank_base_address = output_addr, .page_size = tile_bytes, .data_format = data_format};
+
+    uint32_t end_row = start_row + num_rows_to_process;
+
+    for (uint32_t r = start_row; r < end_row; r++) {
+        cb_wait_front(cb_output_idx, onetile);                // wait until cb has block_size tiles(1 tile)
path: tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/kernels/dataflow/writer_cross_entropy_fw_interleaved_start_id.cpp
user: login: rfurko-tt
body: remove mention about block
created_at: 2025-04-24T18:51:15Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,37 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include "dataflow_api.h"
+
+void kernel_main() {
+    uint32_t runtime_args_counter = 0;
+    uint32_t output_addr = get_arg_val<uint32_t>(runtime_args_counter++);
+    uint32_t num_rows_to_process = get_arg_val<uint32_t>(runtime_args_counter++);
+    uint32_t start_row = get_arg_val<uint32_t>(runtime_args_counter++);
+
+    constexpr uint32_t cb_output_idx = tt::CBIndex::c_11;
+
+    constexpr uint32_t block_size = get_compile_time_arg_val(0);
+    constexpr uint32_t Wt = get_compile_time_arg_val(1);  // number of tiles in inner dimension
+
+    constexpr uint32_t onetile = 1U;
+
+    const uint32_t tile_bytes = get_tile_size(cb_output_idx);
+    const DataFormat data_format = get_dataformat(cb_output_idx);
+
+    const InterleavedAddrGenFast</* is dram */ true> output_addr_generator = {
+        .bank_base_address = output_addr, .page_size = tile_bytes, .data_format = data_format};
+
+    uint32_t end_row = start_row + num_rows_to_process;
+
+    for (uint32_t r = start_row; r < end_row; r++) {
+        cb_wait_front(cb_output_idx, onetile);                // wait until cb has block_size tiles(1 tile)
+        uint32_t l1_read_addr = get_read_ptr(cb_output_idx);  // get the address output buffer
+
+        // write block_size tiles to output buffer
+        noc_async_write_tile(r, output_addr_generator, l1_read_addr);  // write the tile to the output buffer
path: tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/kernels/dataflow/writer_cross_entropy_fw_interleaved_start_id.cpp
user: login: rfurko-tt
body: remove comments
created_at: 2025-04-24T18:52:14Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,287 @@
+
+// SPDX-FileCopyrightText: (c) 2025 Tenstorrent AI ULC
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include <gtest/gtest.h>
+#include <sys/types.h>
+
+#include <cassert>
+#include <core/ttnn_all_includes.hpp>
+#include <cstddef>
+#include <cstdint>
+#include <ttnn/operations/reduction/generic/generic_reductions.hpp>
+
+#include "autograd/auto_context.hpp"
+#include "core/tt_tensor_utils.hpp"
+#include "metal/operations.hpp"
+
+class CrossEntropyForwardTest : public ::testing::Test {
+protected:
+    void SetUp() override {
+        ttml::autograd::ctx().open_device();
+    }
+
+    void TearDown() override {
+        ttml::autograd::ctx().close_device();
+    }
+};
+
+xt::xarray<float> calculate_cross_entropy_loss(const xt::xarray<float>& input, const xt::xarray<uint32_t>& target) {
+    const uint32_t N = target.shape(0);
+    const uint32_t C = 1U;
+    const uint32_t H = target.shape(1);
+    const uint32_t W = 1U;
+    xt::xarray<float> target_inputs = xt::zeros<float>({N, C, H, W});
+
+    for (size_t n = 0; n < N; ++n) {
+        for (size_t h = 0; h < H; ++h) {
+            size_t class_index = target(n, h);
+            target_inputs(n, 0, h, 0) = input(n, 0, h, class_index);
+        }
+    }
+
+    xt::xarray<float> max_input = xt::amax(input, -1, xt::keep_dims);
+    xt::xarray<float> shifted_input = input - max_input;
+    xt::xarray<float> log_exp_sum_test = xt::log(xt::sum(xt::exp(shifted_input), -1, xt::keep_dims));
+    xt::xarray<float> result = -target_inputs + max_input + log_exp_sum_test;
+    return result;
+}
+
+TEST_F(CrossEntropyForwardTest, CrossEntropyForward_Small_Forward) {
+    using namespace ttml;
+
+    const uint32_t N = 1, C = 1, H = 1, W = 8;
+
+    xt::xarray<float> input_tensor = {{{{1.F, 2.F, 3.F, 4.F, 1.F, 2.F, 3.F, 4.F}}}};
+    auto input = core::from_xtensor(input_tensor, &autograd::ctx().get_device());
+    std::cout << "Input Logits:\n";
+    input.print();
+
+    xt::xarray<uint32_t> target_tensor = xt::zeros<uint32_t>({N, H});
+    target_tensor(0, 0) = 1U;
+    auto target = core::from_xtensor<uint32_t, ttnn::DataType::UINT32>(
+        target_tensor, &autograd::ctx().get_device(), ttnn::Layout::ROW_MAJOR);
+    std::cout << "Input Target Indexes:\n";
+    target.print();
+
+    auto result = ttml::metal::cross_entropy_fw(input, target);
+    std::cout << "CrossEntropyForward_Test:\nResult:\n";
+    result.print();
+
+    auto expected_result = calculate_cross_entropy_loss(input_tensor, target_tensor);
+    auto expected_result_print = core::from_xtensor(expected_result, &autograd::ctx().get_device());
+    std::cout << "Expected Result:\n";
+    expected_result_print.print();
+
+    // Check if the result is close to the expected result
+    auto result_xtensor = core::to_xtensor(result);
+    assert((result_xtensor.shape() == expected_result.shape()));
+    EXPECT_TRUE(xt::allclose(result_xtensor, expected_result, 3e-2F, 1e-2F));
+}
+
+TEST_F(CrossEntropyForwardTest, CrossEntropyForward_Negetive_Values) {
+    using namespace ttml;
+
+    const uint32_t N = 1, C = 1, H = 2, W = 4;
+
+    xt::xarray<float> input_tensor = {{{{-100.F, -101.F, -102.F, -103.F}, {-5.01F, -5.02F, -0.3F, -7.F}}}};
+    auto input = core::from_xtensor(input_tensor, &autograd::ctx().get_device());
+    std::cout << "Input Logits:\n";
+    input.print();
+
+    xt::xarray<uint32_t> target_tensor = xt::zeros<uint32_t>({N, H});
+    target_tensor(0, 0) = 0;
+    target_tensor(0, 1) = 2U;
+    auto target = core::from_xtensor<uint32_t, ttnn::DataType::UINT32>(
+        target_tensor, &autograd::ctx().get_device(), ttnn::Layout::ROW_MAJOR);
+    std::cout << "Input Target Indexes:\n";
+    target.print();
+
+    auto result = ttml::metal::cross_entropy_fw(input, target);
+    std::cout << "CrossEntropyForward_Test:\nResult:\n";
+    result.print();
+
+    auto expected_result = calculate_cross_entropy_loss(input_tensor, target_tensor);
+    auto expected_result_print = core::from_xtensor(expected_result, &autograd::ctx().get_device());
+    std::cout << "Expected Result:\n";
+    expected_result_print.print();
+
+    // Check if the result is close to the expected result
+    auto result_xtensor = core::to_xtensor(result);
+    assert((result_xtensor.shape() == expected_result.shape()));
+    EXPECT_TRUE(xt::allclose(result_xtensor, expected_result, 3e-2F, 1e-2F));
+}
+
+TEST_F(CrossEntropyForwardTest, CrossEntropyForward_Batch) {
+    using namespace ttml;
+
+    const uint32_t N = 2U, C = 1U, H = 91U, W = 157U;
+    const auto shape = ttnn::SmallVector<uint32_t>{N, C, H, W};
+
+    std::random_device rd;
+    std::mt19937 gen(rd());  // or fixed seed: std::mt19937 gen(42);
+    xt::xarray<float> input_tensor = xt::random::rand<float>({N, C, H, W}, -10.0F, 10.0F, gen);
+    xt::xarray<uint32_t> target_tensor = xt::zeros<uint32_t>({N, H});
+
+    std::uniform_int_distribution<uint32_t> class_dist(0, W - 1);
+    for (uint32_t n = 0; n < N; ++n) {
+        for (uint32_t h = 0; h < H; ++h) {
+            uint32_t true_class = class_dist(gen);
+            target_tensor(n, h) = true_class;
+        }
+    }
+
+    auto input = core::from_xtensor(input_tensor, &autograd::ctx().get_device());
+    std::cout << "Input Logits:\n";
+    input.print();
+
+    auto target = core::from_xtensor<uint32_t, ttnn::DataType::UINT32>(
+        target_tensor, &autograd::ctx().get_device(), ttnn::Layout::ROW_MAJOR);
+    std::cout << "Input Target Indexes:\n";
+    target.print();
+
+    auto result = ttml::metal::cross_entropy_fw(input, target);
+    std::cout << "CrossEntropyForward_Test:\nResult:\n";
+    result.print();
+
+    auto expected_result = calculate_cross_entropy_loss(input_tensor, target_tensor);
+
+    auto expected_result_print = core::from_xtensor(expected_result, &autograd::ctx().get_device());
+    std::cout << "Expected Result:\n";
+    expected_result_print.print();
+
+    // Check if the result is close to the expected result
+    auto result_xtensor = core::to_xtensor(result);
+    assert((result_xtensor.shape() == expected_result.shape()));
+    EXPECT_TRUE(xt::allclose(result_xtensor, expected_result, 3e-2F, 1e-2F));
+}
+
+TEST_F(CrossEntropyForwardTest, CrossEntropyForward_Large_Batch) {
+    using namespace ttml;
+
+    const uint32_t N = 64U, C = 1U, H = 1017U, W = 1018U;
+    const auto shape = ttnn::SmallVector<uint32_t>{N, C, H, W};
+
+    std::random_device rd;
+    std::mt19937 gen(rd());  // or fixed seed: std::mt19937 gen(42);
+    xt::xarray<float> input_tensor = xt::random::rand<float>({N, C, H, W}, -10.0F, 10.0F, gen);
+    xt::xarray<uint32_t> target_tensor = xt::zeros<uint32_t>({N, H});
+
+    std::uniform_int_distribution<uint32_t> class_dist(0, W - 1);
+    for (uint32_t n = 0; n < N; ++n) {
+        for (uint32_t h = 0; h < H; ++h) {
+            uint32_t true_class = class_dist(gen);
+            target_tensor(n, h) = true_class;
+        }
+    }
+
+    auto input = core::from_xtensor(input_tensor, &autograd::ctx().get_device());
+    std::cout << "Input Logits:\n";
+    input.print();
+
+    auto target = core::from_xtensor<uint32_t, ttnn::DataType::UINT32>(
+        target_tensor, &autograd::ctx().get_device(), ttnn::Layout::ROW_MAJOR);
+    std::cout << "Input Target Indexes:\n";
+    target.print();
+
+    auto result = ttml::metal::cross_entropy_fw(input, target);
+    std::cout << "CrossEntropyForward_Test:\nResult:\n";
+    result.print();
+
+    auto expected_result = calculate_cross_entropy_loss(input_tensor, target_tensor);
+    auto expected_result_print = core::from_xtensor(expected_result, &autograd::ctx().get_device());
+    std::cout << "Expected Result:\n";
+    expected_result_print.print();
+
+    // Check if the result is close to the expected result
+    auto result_xtensor = core::to_xtensor(result);
+    assert((result_xtensor.shape() == expected_result.shape()));
+    EXPECT_TRUE(xt::allclose(result_xtensor, expected_result, 3e-2F, 1e-2F));
+}
+
+TEST_F(CrossEntropyForwardTest, CrossEntropyForward_Large_Forward) {
+    using namespace ttml;
+
+    const uint32_t N = 1U, C = 1U, H = 1U, W = 65536U;
+    const auto shape = ttnn::SmallVector<size_t>{N, C, H, W};
+
+    std::random_device rd;
+    std::mt19937 gen(rd());  // or fixed seed: std::mt19937 gen(42);
+    xt::xarray<float> input_tensor = xt::random::rand<float>({N, C, H, W}, -10.0F, 10.0F, gen);
+    xt::xarray<uint32_t> target_tensor = xt::zeros<uint32_t>({N, H});
+
+    std::uniform_int_distribution<uint32_t> class_dist(0, W - 1);
+    for (uint32_t n = 0; n < N; ++n) {
+        for (uint32_t h = 0; h < H; ++h) {
+            uint32_t true_class = class_dist(gen);
+            target_tensor(n, h) = true_class;
+        }
+    }
+
+    auto input = core::from_xtensor(input_tensor, &autograd::ctx().get_device());
+    std::cout << "Input Logits:\n";
+    input.print();
+
+    auto target = core::from_xtensor<uint32_t, ttnn::DataType::UINT32>(
+        target_tensor, &autograd::ctx().get_device(), ttnn::Layout::ROW_MAJOR);
+    std::cout << "Input Target Indexes:\n";
+    target.print();
+
+    auto result = ttml::metal::cross_entropy_fw(input, target);
+    std::cout << "CrossEntropyForward_Test:\nResult:\n";
+    result.print();
+
+    auto expected_result = calculate_cross_entropy_loss(input_tensor, target_tensor);
+    auto expected_result_print = core::from_xtensor(expected_result, &autograd::ctx().get_device());
+    std::cout << "Expected Result:\n";
+    expected_result_print.print();
+
+    // Check if the result is close to the expected result
+    auto result_xtensor = core::to_xtensor(result);
+    assert((result_xtensor.shape() == expected_result.shape()));
+    EXPECT_TRUE(xt::allclose(result_xtensor, expected_result, 3e-2F, 1e-2F));
+}
+
+TEST_F(CrossEntropyForwardTest, CrossEntropyForward_Huge_Forward) {
+    using namespace ttml;
+
+    const uint32_t N = 64U, C = 1U, H = 32U, W = 128000U;
+    const auto shape = ttnn::SmallVector<size_t>{N, C, H, W};
+
+    std::random_device rd;
+    std::mt19937 gen(rd());  // or fixed seed: std::mt19937 gen(42);
path: tt-train/tests/ops/cross_entropy_fw_op_test.cpp
user: login: rfurko-tt
body: change to fixed seed please
created_at: 2025-04-24T19:09:49Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,474 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include "cross_entropy_fw_program_factory.hpp"
+
+#include <bit>
+#include <cstdint>
+
+#include "cross_entropy_fw_device_operation_types.hpp"
+
+namespace {
+
+constexpr auto kReaderKernelPath =
+    "tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/kernels/dataflow/"
+    "reader_cross_entropy_fw_interleaved_start_id.cpp";
+
+constexpr auto kWriterKernelPath =
+    "tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/kernels/dataflow/"
+    "writer_cross_entropy_fw_interleaved_start_id.cpp";
+
+constexpr auto kComputeKernelPath =
+    "tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/kernels/compute/cross_entropy_fw_kernel.cpp";
+
+// reader runtime args
+constexpr uint32_t kInputBufferIdx = 0;
+constexpr uint32_t kTargetBufferIdx = 1U;
+// writer runtime args
+constexpr uint32_t kOutputBufferIdx = 0;
+
+constexpr auto kInputCbIndex = tt::CBIndex::c_0;
+constexpr auto kTargetCbIndex = tt::CBIndex::c_1;
+constexpr auto kMaskCbIndex = tt::CBIndex::c_2;
+constexpr auto kMaxMaskCbIndex = tt::CBIndex::c_3;
+constexpr auto kScalerCbIndex = tt::CBIndex::c_4;  // used to reduction
+constexpr auto kMaxValueBeforeReductionCbIndex = tt::CBIndex::c_5;
+constexpr auto kMaxValueAfterReductionCbIndex = tt::CBIndex::c_6;
+constexpr auto kExpSumBeforeReductionCbIndex = tt::CBIndex::c_7;
+constexpr auto KExpSumAfterReductionCbIndex = tt::CBIndex::c_8;
+constexpr auto kOutputBeforeReductionCbIndex = tt::CBIndex::c_9;
+constexpr auto kOutputCbIndex = tt::CBIndex::c_10;
+
+constexpr uint32_t kNumMaskTiles = 1U;
+constexpr uint32_t kNumMaxValueTiles = 1U;
+constexpr uint32_t kNumOutputTiles = 1U;
+constexpr uint32_t kMaxValueBeforeReductionTiles = 1U;
+constexpr uint32_t kNumMaxValueAfterReductionTiles = 1U;
+constexpr uint32_t kNumExpSumBeforeReductionTiles = 1U;
+constexpr uint32_t kNumExpSumAfterReductionTiles = 1U;
+constexpr uint32_t kNumScalerTiles = 1U;  // used it to reduction ???
+constexpr uint32_t kNumOutputBeforeReductionTiles = 1U;
+
+const std::string kMaskWDefineKey = "DO_MASK_W";
+
+uint32_t get_block_size(uint32_t num_inner) {
+    const uint32_t max_block_size = 4U;  // 4 is the maximum block size for enabled fp32 dest acc
+    for (uint32_t block_size = max_block_size; block_size > 1; block_size--) {
+        if (num_inner % block_size == 0) {  // if num_inner is divisible by block_size - choose this block_size
+            return block_size;
+        }
+    }
+    return 1U;
+}
+
+}  // namespace
+
+namespace ttml::metal::ops::cross_entropy_fw::device {
+
+/**
+ *   Helper struct to hold references to all kernels we create,
+ *        used during runtime argument setup.
+ */
+struct CrossEntropyForwardKernels {
+    tt::tt_metal::KernelHandle reader;
+    tt::tt_metal::KernelHandle writer;
+    tt::tt_metal::KernelHandle compute_group_1;
+    tt::tt_metal::KernelHandle compute_group_2;
+};
+
+/**
+ *   Create and configure a circular buffer, returning both the configuration and the handle.
+ */
+tt::tt_metal::CBHandle create_circular_buffer(
+    tt::tt_metal::Program& program,
+    const tt::tt_metal::CoreRangeSet& core_ranges,
+    uint32_t cb_index,
+    tt::DataFormat data_format,
+    uint32_t single_tile_size,
+    uint32_t num_tiles) {
+    tt::tt_metal::CircularBufferConfig cb_config =
+        tt::tt_metal::CircularBufferConfig(num_tiles * single_tile_size, {{cb_index, data_format}})
+            .set_page_size(cb_index, single_tile_size);
+
+    auto cb_handle = CreateCircularBuffer(program, core_ranges, cb_config);
+    return cb_handle;
+}
+
+/**
+ *   Create a reader kernel with the given compile-time arguments.
+ */
+tt::tt_metal::KernelHandle create_reader_kernel(
+    tt::tt_metal::Program& program,
+    const tt::tt_metal::CoreRangeSet& core_ranges,
+    const std::vector<uint32_t>& compile_time_args,
+    const std::map<std::string, std::string>& defines,
+    const std::string& kernel_path) {
+    return tt::tt_metal::CreateKernel(
+        program, kernel_path, core_ranges, tt::tt_metal::ReaderDataMovementConfig(compile_time_args, defines));
+}
+
+/**
+ *   Create a writer kernel with the given compile-time arguments.
+ */
+tt::tt_metal::KernelHandle create_writer_kernel(
+    tt::tt_metal::Program& program,
+    const tt::tt_metal::CoreRangeSet& core_ranges,
+    const std::vector<uint32_t>& compile_time_args,
+    const std::map<std::string, std::string>& defines,
+    const std::string& kernel_path) {
+    return tt::tt_metal::CreateKernel(
+        program, kernel_path, core_ranges, tt::tt_metal::WriterDataMovementConfig(compile_time_args, defines));
+}
+
+/**
+ * Create a compute kernel with the given compile-time arguments.
+ */
+tt::tt_metal::KernelHandle create_compute_kernel(
+    tt::tt_metal::Program& program,
+    const tt::tt_metal::CoreRangeSet& core_ranges,
+    const std::vector<uint32_t>& compile_time_args,
+    const std::map<std::string, std::string>& defines,
+    const std::string& kernel_path) {
+    return tt::tt_metal::CreateKernel(
+        program,
+        kernel_path,
+        core_ranges,
+        tt::tt_metal::ComputeConfig{
+            .math_fidelity = MathFidelity::HiFi4,
+            .fp32_dest_acc_en = true,
+            .math_approx_mode = false,
+            .compile_args = compile_time_args,
+            .defines = defines});
+}
+
+/**
+ * Set up the runtime arguments for the 4 relevant kernels (reader, writer, compute G1, compute G2)
+ *        for each core in the grid.
+ */
+void assign_per_core_runtime_args(
+    tt::tt_metal::Program& program,
+    const CrossEntropyForwardKernels& kernels,
+    const tt::tt_metal::Buffer* input_buffer,
+    const tt::tt_metal::Buffer* target_buffer,
+    const tt::tt_metal::Buffer* output_buffer,
+    uint32_t num_cores,
+    uint32_t num_cores_y,
+    uint32_t num_rows_per_core_group_1,
+    uint32_t num_rows_per_core_group_2,
+    const tt::tt_metal::CoreRangeSet& core_group_1,
+    const tt::tt_metal::CoreRangeSet& core_group_2) {
+    for (uint32_t i = 0, num_rows_written = 0; i < num_cores; i++) {
+        CoreCoord core = {i / num_cores_y, i % num_cores_y};
+
+        // Determine how many rows this core will process
+        uint32_t num_rows_per_core = 0;
+        if (core_group_1.contains(core)) {
+            num_rows_per_core = num_rows_per_core_group_1;
+        } else if (core_group_2.contains(core)) {
+            num_rows_per_core = num_rows_per_core_group_2;
+        } else {
+            TT_FATAL(false, "Core not in specified core ranges");
+        }
+
+        // Reader kernel: (input_addr, gamma_addr, number_of_rows, offset_in_rows)
+        SetRuntimeArgs(
+            program,
+            kernels.reader,
+            core,
+            {input_buffer->address(), target_buffer->address(), num_rows_per_core, num_rows_written});
+
+        // Writer kernel: (dst_addr, dst_rms_addr number_of_rows, offset_in_rows)
+        SetRuntimeArgs(
+            program,
+            kernels.writer,
+            core,
+            {output_buffer->address(), /*temporary*/ num_rows_per_core, num_rows_written});
+
+        num_rows_written += num_rows_per_core;
+    }
+}
+
+CrossEntropyForwardProgramFactory::cached_program_t CrossEntropyForwardProgramFactory::create(
+    const operation_attributes_t& operation_attributes,
+    const tensor_args_t& tensor_args,
+    tensor_return_value_t& output) {
+    // -------------------------------------------------------------------------
+    // 1) Setup device, data formats, tile sizes, and compute split
+    // -------------------------------------------------------------------------
+    const auto& input = tensor_args.input;
+    const auto& target = tensor_args.target;
+
+    auto* device = input.device();
+
+    tt::tt_metal::Program program{};
+
+    tt::DataFormat input_data_format = datatype_to_dataformat_converter(input.get_dtype());
+    TT_FATAL(input_data_format == tt::DataFormat::Float16_b, "Input data format must be Float16_b");
+
+    uint32_t bfloat16_single_tile_size_bytes = tt::tt_metal::detail::TileSize(tt::DataFormat::Float16_b);
+    uint32_t float32_single_tile_size_bytes = tt::tt_metal::detail::TileSize(tt::DataFormat::Float32);
+
+    auto padded_tensor_shape = input.get_padded_shape();
+    auto padded_tensor_volume = input.volume();
+
+    TT_FATAL(
+        padded_tensor_volume % tt::constants::TILE_HW == 0, "Padded input tensor volume must be divisible by TILE_HW");
+    TT_FATAL(padded_tensor_shape.rank() == 4U, "Input tensor must be 4D");
+
+    //
+    uint32_t Wt = padded_tensor_shape[-1] / tt::constants::TILE_WIDTH;  // <- number of tiles in inner dimension
+    uint32_t Ht = padded_tensor_shape[-2] / tt::constants::TILE_HEIGHT;
+    uint32_t NC = padded_tensor_shape[0] * padded_tensor_shape[1];
+    uint32_t total_rows_to_process = NC * Ht;
+
+    // get number of free cores
+    auto compute_with_storage_grid_size = device->compute_with_storage_grid_size();
+    uint32_t num_cores_x = compute_with_storage_grid_size.x;
+    uint32_t num_cores_y = compute_with_storage_grid_size.y;
+
+    // get the number of inner dimension(D) - we will caclulate log softmax on this dimension
+    uint32_t num_inner = input.get_logical_shape()[-1];  // (B, 1, S, D) and we take D here
+
+    // mask_w - this mask used to avoid calculation of extra data(data which will be added to create full tile 32x32)??
+    uint32_t mask_w = num_inner % tt::constants::TILE_WIDTH;  // <- width index of first trash value in tile
+
+    // compile arguments
+    uint32_t block_size = get_block_size(Wt);
+
+    auto [num_cores, all_cores, core_group_1, core_group_2, num_rows_per_core_group_1, num_rows_per_core_group_2] =
+        tt::tt_metal::split_work_to_cores(compute_with_storage_grid_size, total_rows_to_process);
+
+    // -------------------------------------------------------------------------
+    // 2) Create and configure circular buffers
+    // -------------------------------------------------------------------------
+
+    // TODO: add calculation of required  memory and check if it fits in L1
+    // then add possibility to kernels to process data by blocks
+    const uint32_t available_L1_in_bytes =
+        device->l1_size_per_core() - device->allocator()->get_base_allocator_addr(tt::tt_metal::HalMemType::L1);
+
+    const uint64_t mask_scaler_maxmask_memory =
+        (kNumMaskTiles + kNumScalerTiles + kNumMaskTiles) * bfloat16_single_tile_size_bytes;
+
+    const uint32_t num_input_tiles = Wt;
+    const uint32_t num_target_tiles = Wt;
+
+    auto data_format = input_data_format;  // tt::DataFormat::Float16_b
+    auto precise_data_format = tt::DataFormat::Float32;
+
+    auto cb_input = create_circular_buffer(
+        program, all_cores, kInputCbIndex, data_format, bfloat16_single_tile_size_bytes, num_input_tiles);
+
+    auto cb_target = create_circular_buffer(
+        program, all_cores, kTargetCbIndex, data_format, bfloat16_single_tile_size_bytes, num_target_tiles);
+
+    auto cb_mask = create_circular_buffer(
+        program, all_cores, kMaskCbIndex, data_format, bfloat16_single_tile_size_bytes, kNumMaskTiles);
+
+    auto cb_max_mask = create_circular_buffer(
+        program, all_cores, kMaxMaskCbIndex, data_format, bfloat16_single_tile_size_bytes, kNumMaskTiles);
+
+    auto cb_scaler = create_circular_buffer(
+        program, all_cores, kScalerCbIndex, data_format, bfloat16_single_tile_size_bytes, kNumScalerTiles);
+
+    auto cb_max_value_before_reduction = create_circular_buffer(
+        program,
+        all_cores,
+        kMaxValueBeforeReductionCbIndex,
+        data_format,
+        bfloat16_single_tile_size_bytes,
+        kMaxValueBeforeReductionTiles);
+
+    auto cb_max_value_after_reduction = create_circular_buffer(
+        program,
+        all_cores,
+        kMaxValueAfterReductionCbIndex,
+        data_format,
+        bfloat16_single_tile_size_bytes,
+        kNumMaxValueAfterReductionTiles);
+
+    auto cb_exp_sum_before_reduction = create_circular_buffer(
+        program,
+        all_cores,
+        kExpSumBeforeReductionCbIndex,
+        data_format,
+        bfloat16_single_tile_size_bytes,
+        kNumExpSumBeforeReductionTiles);
+
+    auto cb_exp_sum_after_refuction = create_circular_buffer(
+        program,
+        all_cores,
+        KExpSumAfterReductionCbIndex,
+        data_format,
+        bfloat16_single_tile_size_bytes,
+        kNumExpSumAfterReductionTiles);
+
+    auto cb_output_before_refuction = create_circular_buffer(
+        program,
+        all_cores,
+        kOutputBeforeReductionCbIndex,
+        data_format,
+        bfloat16_single_tile_size_bytes,
+        kNumOutputBeforeReductionTiles);
+
+    auto cb_output = create_circular_buffer(
+        program,
+        all_cores,
+        kOutputCbIndex,
+        data_format,
+        bfloat16_single_tile_size_bytes,
+        2 * block_size);  // create twice block_size variable
path: tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/cross_entropy_fw_program_factory.cpp
user: login: rfurko-tt
body: you want to have two, while writing one tile you want to start preparing another one
created_at: 2025-04-24T19:13:21Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,133 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include "cross_entropy_fw_device_operation.hpp"
+
+#include "cross_entropy_fw_program_factory.hpp"
+
+namespace ttml::metal::ops::cross_entropy_fw::device {
+
+CrossEntropyForwardDeviceOperation::program_factory_t CrossEntropyForwardDeviceOperation::select_program_factory(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    return CrossEntropyForwardProgramFactory{};
+}
+
+void CrossEntropyForwardDeviceOperation::validate_on_program_cache_hit(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    validate_on_program_cache_miss(args, tensor_args);
+}
+
+void CrossEntropyForwardDeviceOperation::validate_on_program_cache_miss(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    auto check_tensor = [](const ttnn::Tensor& tensor,
+                           const std::string& name,
+                           const tt::tt_metal::Layout required_layout,
+                           const tt::tt_metal::DataType required_dtype) {
+        TT_FATAL(
+            tensor.device()->arch() == tt::ARCH::WORMHOLE_B0,
+            "CrossEntropyForward operation is only supported on Wormhole. Device arch: {}. Tensor name {}",
+            magic_enum::enum_name(tensor.device()->arch()),
+            name);
+
+        TT_FATAL(
+            tensor.storage_type() == tt::tt_metal::StorageType::DEVICE,
+            "CrossEntropyForward operation requires {} to be on Device. Input storage type: {}",
+            name,
+            static_cast<int>(tensor.storage_type()));
+
+        TT_FATAL(
+            tensor.buffer() != nullptr,
+            "Operands to CrossEntropyForward need to be allocated in buffers on the device. Buffer is null. Tensor "
+            "name {}",
+            name);
+
+        TT_FATAL(
+            tensor.get_layout() == required_layout,
+            "CrossEntropyForward operation requires tensor to be in Tile layout. {} tensor layout: {}",
+            name,
+            static_cast<int>(tensor.get_layout()));
+
+        TT_FATAL(
+            tensor.get_dtype() == required_dtype,
+            "CrossEntropyForward operation requires tensor to be of BFLOAT16 data type. {} tensor data type: {}",
+            name,
+            static_cast<int>(tensor.get_dtype()));
+
+        TT_FATAL(
+            tensor.memory_config().memory_layout == ttnn::TensorMemoryLayout::INTERLEAVED,
+            "CrossEntropyForward operation requires Interleaved memory layout. {} "
+            "memory layout: `{}`",
+            name,
+            static_cast<int>(tensor.memory_config().memory_layout));
+    };
+
+    const auto& input_tensor = tensor_args.input;
+    const auto& target_tensor = tensor_args.target;
+    const auto& preallocated_output_tensor = tensor_args.preallocated_output;
+    check_tensor(input_tensor, "Input", tt::tt_metal::Layout::TILE, tt::tt_metal::DataType::BFLOAT16);
+    check_tensor(target_tensor, "Target", tt::tt_metal::Layout::ROW_MAJOR, tt::tt_metal::DataType::UINT32);
+    if (preallocated_output_tensor.has_value()) {
+        check_tensor(
+            preallocated_output_tensor.value(),
+            "Preallocated Output",
+            tt::tt_metal::Layout::TILE,
+            tt::tt_metal::DataType::BFLOAT16);
+    }
+}
+
+CrossEntropyForwardDeviceOperation::spec_return_value_t CrossEntropyForwardDeviceOperation::compute_output_specs(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    if (tensor_args.preallocated_output.has_value()) {
+        return tensor_args.preallocated_output->get_tensor_spec();
+    }
+    auto input_logical_shape = tensor_args.input.get_logical_shape();
+    input_logical_shape[-1] = 1;
+    return ttnn::TensorSpec(
+        ttnn::Shape(input_logical_shape),
+        tt::tt_metal::TensorLayout(
+            tensor_args.input.get_dtype(), tt::tt_metal::Layout::TILE, tensor_args.input.memory_config()));
+}
+
+CrossEntropyForwardDeviceOperation::tensor_return_value_t CrossEntropyForwardDeviceOperation::create_output_tensors(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    tensor_return_value_t output_tensor;
+
+    spec_return_value_t output_specs = compute_output_specs(args, tensor_args);
+
+    if (tensor_args.preallocated_output.has_value()) {
+        output_tensor = tensor_args.preallocated_output.value();
+    } else {
+        output_tensor = create_device_tensor(output_specs, tensor_args.input.device());
+    }
+
+    return output_tensor;
+}
+
+tt::stl::hash::hash_t CrossEntropyForwardDeviceOperation::compute_program_hash(
+    const operation_attributes_t& args, const tensor_args_t& tensor_args) {
+    const auto& input_tensor = tensor_args.input;
+    const auto& input_logical_shape = input_tensor.get_logical_shape();
+    auto program_factory = select_program_factory(args, tensor_args);
+    tt::tt_metal::operation::Hash hash = tt::tt_metal::operation::hash_operation<CrossEntropyForwardDeviceOperation>(
+        args, program_factory.index(), input_tensor.dtype(), input_logical_shape);
+
+    return hash;
+}
+
+std::
+    tuple<CrossEntropyForwardDeviceOperation::operation_attributes_t, CrossEntropyForwardDeviceOperation::tensor_args_t>
path: tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/cross_entropy_fw_device_operation.cpp
user: login: vmelnykovTT
body: This formatting happens automatically when I save or run clang-format, so I’m not manually formatting it that way
created_at: 2025-04-25T13:07:49Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,508 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include "cross_entropy_fw_program_factory.hpp"
+
+#include <bit>
+#include <cstdint>
+#include <tt-metalium/buffer.hpp>
+
+#include "cross_entropy_fw_device_operation_types.hpp"
+
+namespace {
+
+constexpr auto kReaderKernelPath =
+    "tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/kernels/dataflow/"
+    "reader_cross_entropy_fw_interleaved_start_id.cpp";
+
+constexpr auto kWriterKernelPath =
+    "tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/kernels/dataflow/"
+    "writer_cross_entropy_fw_interleaved_start_id.cpp";
+
+constexpr auto kComputeKernelPath =
+    "tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/kernels/compute/cross_entropy_fw_kernel.cpp";
+
+// reader runtime args
+constexpr uint32_t kInputBufferIdx = 0;
+constexpr uint32_t kTargetBufferIdx = 1U;
+// writer runtime args
+constexpr uint32_t kOutputBufferIdx = 0;
+
+constexpr auto kInputCbIndex = tt::CBIndex::c_0;
+constexpr auto kTargetCbIndex = tt::CBIndex::c_1;
+constexpr auto kMaskCbIndex = tt::CBIndex::c_2;
+constexpr auto kMaxMaskCbIndex = tt::CBIndex::c_3;
+constexpr auto kScalerCbIndex = tt::CBIndex::c_4;  // used to reduction
+constexpr auto kInputTileCbIndex = tt::CBIndex::c_5;
+constexpr auto kTargetLogitsCbIndex = tt::CBIndex::c_6;
+constexpr auto kMaxValueBeforeReductionCbIndex = tt::CBIndex::c_7;
+constexpr auto kMaxValueAfterReductionCbIndex = tt::CBIndex::c_8;
+constexpr auto kExpSumBeforeReductionCbIndex = tt::CBIndex::c_9;
+constexpr auto KExpSumAfterReductionCbIndex = tt::CBIndex::c_10;
+constexpr auto kOutputCbIndex = tt::CBIndex::c_11;
+
+constexpr uint32_t kNumTargetIndexesTiles = 1U;
+constexpr uint32_t kNumMaskTiles = 1U;
+constexpr uint32_t kNumMaxValueTiles = 1U;
+constexpr uint32_t kNumInputTilesByIndx = 1U;
+constexpr uint32_t kMaxValueBeforeReductionTiles = 1U;
+constexpr uint32_t kNumTargetLogitsTiles = 1U;
+constexpr uint32_t kNumMaxValueAfterReductionTiles = 1U;
+constexpr uint32_t kNumExpSumBeforeReductionTiles = 1U;
+constexpr uint32_t kNumExpSumAfterReductionTiles = 1U;
+constexpr uint32_t kNumScalerTiles = 1U;  // used it to reduction ???
+constexpr uint32_t kNumOutputBeforeReductionTiles = 1U;
+constexpr uint32_t kNumOutputTiles = 1U;
+
+constexpr uint32_t kPageElementsNumber = 32U;
+
+const std::string kMaskWDefineKey = "DO_MASK_W";
+const std::string kEverythingFitsInL1DefineKey = "EVERYTHING_FITS_IN_L1";
+
+uint32_t get_block_size(uint32_t num_inner) {
+    const uint32_t max_block_size = 4U;  // 4 is the maximum block size for enabled fp32 dest acc
+    for (uint32_t block_size = max_block_size; block_size > 1; block_size--) {
+        if (num_inner % block_size == 0) {  // if num_inner is divisible by block_size - choose this block_size
+            return block_size;
+        }
+    }
+    return 1U;
+}
+
+}  // namespace
+
+namespace ttml::metal::ops::cross_entropy_fw::device {
+
+/**
+ *   Helper struct to hold references to all kernels we create,
+ *        used during runtime argument setup.
+ */
+struct CrossEntropyForwardKernels {
+    tt::tt_metal::KernelHandle reader;
+    tt::tt_metal::KernelHandle writer;
+    tt::tt_metal::KernelHandle compute_group_1;
+    tt::tt_metal::KernelHandle compute_group_2;
+};
+
+/**
+ *   Create and configure a circular buffer, returning both the configuration and the handle.
+ */
+tt::tt_metal::CBHandle create_circular_buffer(
+    tt::tt_metal::Program& program,
+    const tt::tt_metal::CoreRangeSet& core_ranges,
+    uint32_t cb_index,
+    tt::DataFormat data_format,
+    uint32_t single_tile_size,
+    uint32_t num_tiles) {
+    tt::tt_metal::CircularBufferConfig cb_config =
+        tt::tt_metal::CircularBufferConfig(num_tiles * single_tile_size, {{cb_index, data_format}})
+            .set_page_size(cb_index, single_tile_size);
+
+    auto cb_handle = CreateCircularBuffer(program, core_ranges, cb_config);
+    return cb_handle;
+}
+
+/**
+ *   Create a reader kernel with the given compile-time arguments.
+ */
+tt::tt_metal::KernelHandle create_reader_kernel(
+    tt::tt_metal::Program& program,
+    const tt::tt_metal::CoreRangeSet& core_ranges,
+    const std::vector<uint32_t>& compile_time_args,
+    const std::map<std::string, std::string>& defines,
+    const std::string& kernel_path) {
+    return tt::tt_metal::CreateKernel(
+        program, kernel_path, core_ranges, tt::tt_metal::ReaderDataMovementConfig(compile_time_args, defines));
+}
+
+/**
+ *   Create a writer kernel with the given compile-time arguments.
+ */
+tt::tt_metal::KernelHandle create_writer_kernel(
+    tt::tt_metal::Program& program,
+    const tt::tt_metal::CoreRangeSet& core_ranges,
+    const std::vector<uint32_t>& compile_time_args,
+    const std::map<std::string, std::string>& defines,
+    const std::string& kernel_path) {
+    return tt::tt_metal::CreateKernel(
+        program, kernel_path, core_ranges, tt::tt_metal::WriterDataMovementConfig(compile_time_args, defines));
+}
+
+/**
+ * Create a compute kernel with the given compile-time arguments.
+ */
+tt::tt_metal::KernelHandle create_compute_kernel(
+    tt::tt_metal::Program& program,
+    const tt::tt_metal::CoreRangeSet& core_ranges,
+    const std::vector<uint32_t>& compile_time_args,
+    const std::map<std::string, std::string>& defines,
+    const std::string& kernel_path) {
+    return tt::tt_metal::CreateKernel(
+        program,
+        kernel_path,
+        core_ranges,
+        tt::tt_metal::ComputeConfig{
+            .math_fidelity = MathFidelity::HiFi4,
+            .fp32_dest_acc_en = true,
+            .math_approx_mode = false,
+            .compile_args = compile_time_args,
+            .defines = defines});
+}
+
+/**
+ * Set up the runtime arguments for the 4 relevant kernels (reader, writer, compute G1, compute G2)
+ *        for each core in the grid.
+ */
+void assign_per_core_runtime_args(
+    tt::tt_metal::Program& program,
+    const CrossEntropyForwardKernels& kernels,
+    const tt::tt_metal::Buffer* input_buffer,
+    const tt::tt_metal::Buffer* target_buffer,
+    const tt::tt_metal::Buffer* output_buffer,
+
+    uint32_t num_cores,
+    uint32_t num_cores_y,
+    uint32_t num_rows_per_core_group_1,
+    uint32_t num_rows_per_core_group_2,
+    const tt::tt_metal::CoreRangeSet& core_group_1,
+    const tt::tt_metal::CoreRangeSet& core_group_2) {
+    for (uint32_t i = 0, num_rows_written = 0; i < num_cores; i++) {
+        CoreCoord core = {i / num_cores_y, i % num_cores_y};
+
+        // Determine how many rows this core will process
+        uint32_t num_rows_per_core = 0;
+        if (core_group_1.contains(core)) {
+            num_rows_per_core = num_rows_per_core_group_1;
+        } else if (core_group_2.contains(core)) {
+            num_rows_per_core = num_rows_per_core_group_2;
+        } else {
+            TT_FATAL(false, "Core not in specified core ranges");
+        }
+
+        // Reader kernel: (input_addr, gamma_addr, number_of_rows, offset_in_rows)
+        SetRuntimeArgs(
+            program,
+            kernels.reader,
+            core,
+            {input_buffer->address(), target_buffer->address(), num_rows_per_core, num_rows_written});
+
+        // Writer kernel: (dst_addr, dst_rms_addr number_of_rows, offset_in_rows)
+        SetRuntimeArgs(
+            program,
+            kernels.writer,
+            core,
+            {output_buffer->address(), /*temporary*/ num_rows_per_core, num_rows_written});
+
+        num_rows_written += num_rows_per_core;
+    }
+}
+
+CrossEntropyForwardProgramFactory::cached_program_t CrossEntropyForwardProgramFactory::create(
+    const operation_attributes_t& operation_attributes,
+    const tensor_args_t& tensor_args,
+    tensor_return_value_t& output) {
+    // -------------------------------------------------------------------------
+    // 1) Setup device, data formats, tile sizes, and compute split
+    // -------------------------------------------------------------------------
+    const auto& input = tensor_args.input;
+    const auto& target = tensor_args.target;
+
+    auto* device = input.device();
+
+    tt::tt_metal::Program program{};
+
+    tt::DataFormat input_data_format = datatype_to_dataformat_converter(input.get_dtype());
+    TT_FATAL(input_data_format == tt::DataFormat::Float16_b, "Input data format must be Float16_b");
+
+    uint32_t bfloat16_single_tile_size_bytes = tt::tt_metal::detail::TileSize(tt::DataFormat::Float16_b);
+    uint32_t float32_single_tile_size_bytes = tt::tt_metal::detail::TileSize(tt::DataFormat::Float32);
+
+    auto padded_tensor_shape = input.get_padded_shape();
+    auto padded_tensor_volume = input.volume();
+
+    TT_FATAL(
+        padded_tensor_volume % tt::constants::TILE_HW == 0, "Padded input tensor volume must be divisible by TILE_HW");
+    TT_FATAL(padded_tensor_shape.rank() == 4U, "Input tensor must be 4D");
+
+    //
+    uint32_t Wt = padded_tensor_shape[-1] / tt::constants::TILE_WIDTH;  // <- number of tiles in inner dimension
+    uint32_t Ht = padded_tensor_shape[-2] / tt::constants::TILE_HEIGHT;
+    uint32_t NC = padded_tensor_shape[0] * padded_tensor_shape[1];
+    uint32_t total_rows_to_process = NC * Ht;
+
+    // get size of target indexes inner dimension
+    uint32_t target_indexes_inner_dim_size = target.get_logical_shape()[-1] * target.element_size();
+    // read target indexes by pages(32 indexes in page)
+    uint32_t uint32_read_page_size = tt::datum_size(tt::DataFormat::UInt32) * kPageElementsNumber;
+    // tiled height of target indexes
+    uint32_t tiled_H =
path: tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/cross_entropy_fw_program_factory.cpp
user: login: vmelnykovTT
body: I don't clearly understand what do you mean?
created_at: 2025-04-25T16:09:49Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,508 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#include "cross_entropy_fw_program_factory.hpp"
+
+#include <bit>
+#include <cstdint>
+#include <tt-metalium/buffer.hpp>
+
+#include "cross_entropy_fw_device_operation_types.hpp"
+
+namespace {
+
+constexpr auto kReaderKernelPath =
+    "tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/kernels/dataflow/"
+    "reader_cross_entropy_fw_interleaved_start_id.cpp";
+
+constexpr auto kWriterKernelPath =
+    "tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/kernels/dataflow/"
+    "writer_cross_entropy_fw_interleaved_start_id.cpp";
+
+constexpr auto kComputeKernelPath =
+    "tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/kernels/compute/cross_entropy_fw_kernel.cpp";
+
+// reader runtime args
+constexpr uint32_t kInputBufferIdx = 0;
+constexpr uint32_t kTargetBufferIdx = 1U;
+// writer runtime args
+constexpr uint32_t kOutputBufferIdx = 0;
+
+constexpr auto kInputCbIndex = tt::CBIndex::c_0;
+constexpr auto kTargetCbIndex = tt::CBIndex::c_1;
+constexpr auto kMaskCbIndex = tt::CBIndex::c_2;
+constexpr auto kMaxMaskCbIndex = tt::CBIndex::c_3;
+constexpr auto kScalerCbIndex = tt::CBIndex::c_4;  // used to reduction
+constexpr auto kInputTileCbIndex = tt::CBIndex::c_5;
+constexpr auto kTargetLogitsCbIndex = tt::CBIndex::c_6;
+constexpr auto kMaxValueBeforeReductionCbIndex = tt::CBIndex::c_7;
+constexpr auto kMaxValueAfterReductionCbIndex = tt::CBIndex::c_8;
+constexpr auto kExpSumBeforeReductionCbIndex = tt::CBIndex::c_9;
+constexpr auto KExpSumAfterReductionCbIndex = tt::CBIndex::c_10;
+constexpr auto kOutputCbIndex = tt::CBIndex::c_11;
+
+constexpr uint32_t kNumTargetIndexesTiles = 1U;
+constexpr uint32_t kNumMaskTiles = 1U;
+constexpr uint32_t kNumMaxValueTiles = 1U;
+constexpr uint32_t kNumInputTilesByIndx = 1U;
+constexpr uint32_t kMaxValueBeforeReductionTiles = 1U;
+constexpr uint32_t kNumTargetLogitsTiles = 1U;
+constexpr uint32_t kNumMaxValueAfterReductionTiles = 1U;
+constexpr uint32_t kNumExpSumBeforeReductionTiles = 1U;
+constexpr uint32_t kNumExpSumAfterReductionTiles = 1U;
+constexpr uint32_t kNumScalerTiles = 1U;  // used it to reduction ???
+constexpr uint32_t kNumOutputBeforeReductionTiles = 1U;
+constexpr uint32_t kNumOutputTiles = 1U;
+
+constexpr uint32_t kPageElementsNumber = 32U;
+
+const std::string kMaskWDefineKey = "DO_MASK_W";
+const std::string kEverythingFitsInL1DefineKey = "EVERYTHING_FITS_IN_L1";
+
+uint32_t get_block_size(uint32_t num_inner) {
+    const uint32_t max_block_size = 4U;  // 4 is the maximum block size for enabled fp32 dest acc
+    for (uint32_t block_size = max_block_size; block_size > 1; block_size--) {
+        if (num_inner % block_size == 0) {  // if num_inner is divisible by block_size - choose this block_size
+            return block_size;
+        }
+    }
+    return 1U;
+}
+
+}  // namespace
+
+namespace ttml::metal::ops::cross_entropy_fw::device {
+
+/**
+ *   Helper struct to hold references to all kernels we create,
+ *        used during runtime argument setup.
+ */
+struct CrossEntropyForwardKernels {
+    tt::tt_metal::KernelHandle reader;
+    tt::tt_metal::KernelHandle writer;
+    tt::tt_metal::KernelHandle compute_group_1;
+    tt::tt_metal::KernelHandle compute_group_2;
+};
+
+/**
+ *   Create and configure a circular buffer, returning both the configuration and the handle.
+ */
+tt::tt_metal::CBHandle create_circular_buffer(
+    tt::tt_metal::Program& program,
+    const tt::tt_metal::CoreRangeSet& core_ranges,
+    uint32_t cb_index,
+    tt::DataFormat data_format,
+    uint32_t single_tile_size,
+    uint32_t num_tiles) {
+    tt::tt_metal::CircularBufferConfig cb_config =
+        tt::tt_metal::CircularBufferConfig(num_tiles * single_tile_size, {{cb_index, data_format}})
+            .set_page_size(cb_index, single_tile_size);
+
+    auto cb_handle = CreateCircularBuffer(program, core_ranges, cb_config);
+    return cb_handle;
+}
+
+/**
+ *   Create a reader kernel with the given compile-time arguments.
+ */
+tt::tt_metal::KernelHandle create_reader_kernel(
+    tt::tt_metal::Program& program,
+    const tt::tt_metal::CoreRangeSet& core_ranges,
+    const std::vector<uint32_t>& compile_time_args,
+    const std::map<std::string, std::string>& defines,
+    const std::string& kernel_path) {
+    return tt::tt_metal::CreateKernel(
+        program, kernel_path, core_ranges, tt::tt_metal::ReaderDataMovementConfig(compile_time_args, defines));
+}
+
+/**
+ *   Create a writer kernel with the given compile-time arguments.
+ */
+tt::tt_metal::KernelHandle create_writer_kernel(
+    tt::tt_metal::Program& program,
+    const tt::tt_metal::CoreRangeSet& core_ranges,
+    const std::vector<uint32_t>& compile_time_args,
+    const std::map<std::string, std::string>& defines,
+    const std::string& kernel_path) {
+    return tt::tt_metal::CreateKernel(
+        program, kernel_path, core_ranges, tt::tt_metal::WriterDataMovementConfig(compile_time_args, defines));
+}
+
+/**
+ * Create a compute kernel with the given compile-time arguments.
+ */
+tt::tt_metal::KernelHandle create_compute_kernel(
+    tt::tt_metal::Program& program,
+    const tt::tt_metal::CoreRangeSet& core_ranges,
+    const std::vector<uint32_t>& compile_time_args,
+    const std::map<std::string, std::string>& defines,
+    const std::string& kernel_path) {
+    return tt::tt_metal::CreateKernel(
+        program,
+        kernel_path,
+        core_ranges,
+        tt::tt_metal::ComputeConfig{
+            .math_fidelity = MathFidelity::HiFi4,
+            .fp32_dest_acc_en = true,
+            .math_approx_mode = false,
+            .compile_args = compile_time_args,
+            .defines = defines});
+}
+
+/**
+ * Set up the runtime arguments for the 4 relevant kernels (reader, writer, compute G1, compute G2)
+ *        for each core in the grid.
+ */
+void assign_per_core_runtime_args(
+    tt::tt_metal::Program& program,
+    const CrossEntropyForwardKernels& kernels,
+    const tt::tt_metal::Buffer* input_buffer,
+    const tt::tt_metal::Buffer* target_buffer,
+    const tt::tt_metal::Buffer* output_buffer,
+
+    uint32_t num_cores,
+    uint32_t num_cores_y,
+    uint32_t num_rows_per_core_group_1,
+    uint32_t num_rows_per_core_group_2,
+    const tt::tt_metal::CoreRangeSet& core_group_1,
+    const tt::tt_metal::CoreRangeSet& core_group_2) {
+    for (uint32_t i = 0, num_rows_written = 0; i < num_cores; i++) {
+        CoreCoord core = {i / num_cores_y, i % num_cores_y};
+
+        // Determine how many rows this core will process
+        uint32_t num_rows_per_core = 0;
+        if (core_group_1.contains(core)) {
+            num_rows_per_core = num_rows_per_core_group_1;
+        } else if (core_group_2.contains(core)) {
+            num_rows_per_core = num_rows_per_core_group_2;
+        } else {
+            TT_FATAL(false, "Core not in specified core ranges");
+        }
+
+        // Reader kernel: (input_addr, gamma_addr, number_of_rows, offset_in_rows)
+        SetRuntimeArgs(
+            program,
+            kernels.reader,
+            core,
+            {input_buffer->address(), target_buffer->address(), num_rows_per_core, num_rows_written});
+
+        // Writer kernel: (dst_addr, dst_rms_addr number_of_rows, offset_in_rows)
+        SetRuntimeArgs(
+            program,
+            kernels.writer,
+            core,
+            {output_buffer->address(), /*temporary*/ num_rows_per_core, num_rows_written});
+
+        num_rows_written += num_rows_per_core;
+    }
+}
+
+CrossEntropyForwardProgramFactory::cached_program_t CrossEntropyForwardProgramFactory::create(
+    const operation_attributes_t& operation_attributes,
+    const tensor_args_t& tensor_args,
+    tensor_return_value_t& output) {
+    // -------------------------------------------------------------------------
+    // 1) Setup device, data formats, tile sizes, and compute split
+    // -------------------------------------------------------------------------
+    const auto& input = tensor_args.input;
+    const auto& target = tensor_args.target;
+
+    auto* device = input.device();
+
+    tt::tt_metal::Program program{};
+
+    tt::DataFormat input_data_format = datatype_to_dataformat_converter(input.get_dtype());
+    TT_FATAL(input_data_format == tt::DataFormat::Float16_b, "Input data format must be Float16_b");
+
+    uint32_t bfloat16_single_tile_size_bytes = tt::tt_metal::detail::TileSize(tt::DataFormat::Float16_b);
+    uint32_t float32_single_tile_size_bytes = tt::tt_metal::detail::TileSize(tt::DataFormat::Float32);
+
+    auto padded_tensor_shape = input.get_padded_shape();
+    auto padded_tensor_volume = input.volume();
+
+    TT_FATAL(
+        padded_tensor_volume % tt::constants::TILE_HW == 0, "Padded input tensor volume must be divisible by TILE_HW");
+    TT_FATAL(padded_tensor_shape.rank() == 4U, "Input tensor must be 4D");
+
+    //
+    uint32_t Wt = padded_tensor_shape[-1] / tt::constants::TILE_WIDTH;  // <- number of tiles in inner dimension
+    uint32_t Ht = padded_tensor_shape[-2] / tt::constants::TILE_HEIGHT;
+    uint32_t NC = padded_tensor_shape[0] * padded_tensor_shape[1];
+    uint32_t total_rows_to_process = NC * Ht;
+
+    // get size of target indexes inner dimension
+    uint32_t target_indexes_inner_dim_size = target.get_logical_shape()[-1] * target.element_size();
+    // read target indexes by pages(32 indexes in page)
+    uint32_t uint32_read_page_size = tt::datum_size(tt::DataFormat::UInt32) * kPageElementsNumber;
+    // tiled height of target indexes
+    uint32_t tiled_H =
path: tt-train/sources/ttml/metal/ops/cross_entropy_fw/device/cross_entropy_fw_program_factory.cpp
user: login: romaf5
body: Ht = tiled_H
created_at: 2025-04-28T18:00:28Z
author_association: NONE

PR comments:
user: login: rfurko-tt
created_at: 2025-04-24T19:12:49Z
body: Please also update PR description
author_association: CONTRIBUTOR
user: login: rfurko-tt
created_at: 2025-04-24T19:46:37Z
body: it would be also great to update cross entropy operation and MNIST demo 
author_association: CONTRIBUTOR