Repository: tenstorrent/tt-metal
Pull request #22333: Add zero padding for matmul ops
Author: login: nsorabaTT
Created: 2025-05-20T07:36:35Z | Merged: 2025-05-28T14:22:50Z

### Ticket
https://github.com/tenstorrent/tt-metal/issues/22012

### Problem description
Matmul family of ops operate on tile format. When a tensor is represented in tile format, the dimensions are rounded up to the nearest tile dimension. Depending on the tile size, this can add upto 31 extra entries in each dimension.

For example, if the K dimension is 161, and tile width is 16, then the effective size will be 176, adding 15 extra elements in K dimension.

These extra elements should be ignore for matmul computation, specifically on the K dimension as it can add additional values to the scalar dot product.

### What's changed
This PR:
1. Adds a helper function that pads a tile with a given value. This takes pad width, data type, tile size and face size into consideration.
2. All matmul kernels zero out the padded values when reading _input A_ (first operand) to matmul along K dimension.
3. Adds a unit test to cover this scenario.

Note: It is sufficient to pad the first operand along the inner dimension for matmul, since the resulting product is assured to be zero.

<img width="982" alt="image" src="https://github.com/user-attachments/assets/d2570cfd-9170-429f-975b-69df7ce32c7e" />


### Checklist
- [x] [All post commit](https://github.com/tenstorrent/tt-metal/actions/workflows/all-post-commit-workflows.yaml) Passing: https://github.com/tenstorrent/tt-metal/actions/runs/15214400424
- [x] [Blackhole Post commit](https://github.com/tenstorrent/tt-metal/actions/workflows/blackhole-post-commit.yaml) Passing: https://github.com/tenstorrent/tt-metal/actions/runs/15214402343
- [x] [Model regression](https://github.com/tenstorrent/tt-metal/actions/workflows/perf-models.yaml) Passing: https://github.com/tenstorrent/tt-metal/actions/runs/15215037968/job/42798776645
- [x] [Device performance regression](https://github.com/tenstorrent/tt-metal/actions/workflows/perf-device-models.yaml) Passing: https://github.com/tenstorrent/tt-metal/actions/runs/15215226758
- [ ] **(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)
- [ ] New/Existing tests provide coverage for changes


Diff:
diff --git a/tests/ttnn/unit_tests/operations/matmul/test_matmul.py b/tests/ttnn/unit_tests/operations/matmul/test_matmul.py
index 4b626d68dd75..db1e2b4bf421 100644
--- a/tests/ttnn/unit_tests/operations/matmul/test_matmul.py
+++ b/tests/ttnn/unit_tests/operations/matmul/test_matmul.py
@@ -2227,3 +2227,157 @@ def test_sharded_matmul_with_multiple_out_block_values(device, out_block_h, out_
     )
     output_tensor = ttnn.to_torch(output_tensor)
     assert_with_pcc(torch_output_tensor, output_tensor, pcc=pcc)
+
+
+@pytest.mark.parametrize("input_b_value", [2.0])
+@pytest.mark.parametrize("input_a_value", [4.0])
+@pytest.mark.parametrize(
+    "input_a_shape,input_b_shape,input_a_reshape,input_b_reshape",
+    [
+        ((32, 96), (96, 32), (32, 96), (96, 32)),  # No padding introduced
+        ((32, 96), (96, 32), (1, 90), (90, 16)),  # Padding introduced in M,K and N dimensions, 1 face padded
+        ((32, 96), (96, 32), (1, 65), (65, 16)),  # Padding introduced in M,K and N dimensions, 2 faces padded
+    ],
+)
+@pytest.mark.parametrize(
+    "program_config,input_a_memory_config,input_b_memory_config,output_memory_config",
+    [
+        # Uses reader_bmm_tile_layout_in0.cpp
+        (
+            ttnn.MatmulMultiCoreReuseProgramConfig(
+                compute_with_storage_grid_size=(3, 1),
+                in0_block_w=1,
+                out_subblock_h=1,
+                out_subblock_w=1,
+                per_core_M=1,
+                per_core_N=1,
+            ),
+            None,
+            None,
+            None,
+        ),
+        # Uses reader_bmm_tile_layout_in0_sender_padding.cpp
+        (
+            ttnn.MatmulMultiCoreReuseMultiCast1DProgramConfig(
+                compute_with_storage_grid_size=(3, 1),
+                in0_block_w=1,
+                out_subblock_h=1,
+                out_subblock_w=1,
+                per_core_M=1,
+                per_core_N=1,
+                fuse_batch=True,
+                fused_activation=None,
+                mcast_in0=True,
+            ),
+            None,
+            None,
+            None,
+        ),
+        # Uses reader_bmm_tile_layout_in0_sender_receiver_padding_block_sharded.cpp
+        (
+            ttnn.MatmulMultiCoreReuseMultiCast1DProgramConfig(
+                compute_with_storage_grid_size=(3, 1),
+                in0_block_w=1,
+                out_subblock_h=1,
+                out_subblock_w=1,
+                per_core_M=1,
+                per_core_N=1,
+                fuse_batch=True,
+                fused_activation=None,
+                mcast_in0=True,
+            ),
+            ttnn.MemoryConfig(
+                memory_layout=ttnn.TensorMemoryLayout.WIDTH_SHARDED,
+                buffer_type=ttnn.BufferType.L1,
+                shard_spec=ttnn.ShardSpec(
+                    ttnn.CoreRangeSet({ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(2, 0))}),
+                    (32, 96),
+                    ttnn.ShardOrientation.ROW_MAJOR,
+                ),
+            ),
+            None,
+            None,
+        ),
+        # Uses reader_bmm_tile_layout_in0_sender_dram_sharded.cpp
+        (
+            ttnn.MatmulMultiCoreReuseMultiCastDRAMShardedProgramConfig(
+                in0_block_w=1,
+                per_core_M=1,
+                per_core_N=1,
+                fused_activation=None,
+            ),
+            ttnn.MemoryConfig(
+                memory_layout=ttnn.TensorMemoryLayout.WIDTH_SHARDED,
+                buffer_type=ttnn.BufferType.L1,
+                shard_spec=ttnn.ShardSpec(
+                    ttnn.CoreRangeSet({ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(2, 0))}),
+                    (32, 32),
+                    ttnn.ShardOrientation.ROW_MAJOR,
+                ),
+            ),
+            ttnn.MemoryConfig(
+                memory_layout=ttnn.TensorMemoryLayout.WIDTH_SHARDED,
+                buffer_type=ttnn.BufferType.DRAM,
+                shard_spec=ttnn.ShardSpec(
+                    ttnn.CoreRangeSet({ttnn.CoreRange(ttnn.CoreCoord(0, 0), ttnn.CoreCoord(2, 0))}),
+                    (96, 32),
+                    ttnn.ShardOrientation.ROW_MAJOR,
+                ),
+            ),
+            ttnn.MemoryConfig(
+                memory_layout=ttnn.TensorMemoryLayout.WIDTH_SHARDED,
+                buffer_type=ttnn.BufferType.L1,
+            ),
+        ),
+    ],
+)
+def test_matmul_padding(
+    device,
+    input_a_shape,
+    input_b_shape,
+    input_a_value,
+    input_b_value,
+    input_a_reshape,
+    input_b_reshape,
+    program_config,
+    input_a_memory_config,
+    input_b_memory_config,
+    output_memory_config,
+):
+    torch.manual_seed(0)
+
+    # Create input tensors with specified shapes and values
+    input_a = torch.full(input_a_shape, input_a_value, dtype=torch.float32)
+    input_b = torch.full(input_b_shape, input_b_value, dtype=torch.float32)
+
+    # Reshaped tensors for matmul
+    input_a_torch = torch.full(input_a_reshape, input_a_value, dtype=torch.float32)
+    input_b_torch = torch.full(input_b_reshape, input_b_value, dtype=torch.float32)
+
+    # Compute golden output
+    golden_output = torch.matmul(input_a_torch, input_b_torch)
+
+    # Convert to ttnn tensors
+    input_a_ttnn = ttnn.from_torch(
+        input_a, dtype=ttnn.float32, layout=ttnn.TILE_LAYOUT, device=device, memory_config=input_a_memory_config
+    )
+    input_b_ttnn = ttnn.from_torch(
+        input_b, dtype=ttnn.float32, layout=ttnn.TILE_LAYOUT, device=device, memory_config=input_b_memory_config
+    )
+
+    # Reshape ttnn tensors
+    input_a_reshaped_ttnn = ttnn.reshape(input_a_ttnn, input_a_reshape, padded_shape=input_a_shape)
+    input_b_reshaped_ttnn = ttnn.reshape(input_b_ttnn, input_b_reshape, padded_shape=input_b_shape)
+
+    # Compute matmul
+    for _ in range(11):
+        output_ttnn = ttnn.matmul(
+            input_a_reshaped_ttnn,
+            input_b_reshaped_ttnn,
+            program_config=program_config,
+            memory_config=output_memory_config,
+        )
+    output = ttnn.to_torch(output_ttnn)
+
+    # Verify values match with high precision
+    assert torch.allclose(golden_output, output, atol=1e-6)
diff --git a/ttnn/cpp/ttnn/operations/matmul/device/kernels/dataflow/pad_tile.hpp b/ttnn/cpp/ttnn/operations/matmul/device/kernels/dataflow/pad_tile.hpp
new file mode 100644
index 000000000000..60eca81b5a3c
--- /dev/null
+++ b/ttnn/cpp/ttnn/operations/matmul/device/kernels/dataflow/pad_tile.hpp
@@ -0,0 +1,180 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent AI ULC
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#pragma once
+
+#include <tt-metalium/constants.hpp>
+
+using namespace tt::constants;
+
+/**
+ * @brief Pads a face within a tile.
+ *
+ * This function handles padding for a face within a tile. It calculates the appropriate
+ * padding for the face based on the unpadded dimensions and applies it.
+
+    Case 1: All elements in face are padded (num_elements_unpadded_w <= face_w_offset)
+
+    num_elements_unpadded_w
+            v
+    +----------------+----------------+
+    | x x x          |                |
+    | x x x          |     Face       |
+    | x x x          |                |
+    +----------------+----------------+
+    ^                ^
+                face_w_offset
+
+    Case 2: All elements are unpadded (num_elements_unpadded_w >= face_w_offset + FACE_WIDTH)
+    +----------------+----------------+
+    | x x x x x x x x| x x            |
+    | x x Face  x x x| x x            |
+    | x x x x x x x x| x x            |
+    +----------------+----------------+
+    ^                ^
+    face_w_offset    face_w_offset + FACE_WIDTH
+
+    Case 3: Some elements are padded (face_w_offset < num_elements_unpadded_w < face_w_offset + FACE_WIDTH)
+            num_elements_unpadded_w
+                          v
+    +----------------+----------------+
+    | x x x x x x x x| x x            |
+    | x x x x x x x x| x x  Face      |
+    | x x x x x x x x| x x            |
+    +----------------+----------------+
+    ^                ^                ^
+                face_w_offset         face_w_offset + FACE_WIDTH
+
+ * @tparam T The type of the elements in the tile
+ * @tparam num_elements_unpadded_w Number of unpadded elements in the width dimension
+ * @tparam num_elements_unpadded_h Number of unpadded elements in the height dimension
+ * @tparam num_faces_w Number of faces in the width dimension
+ * @tparam num_faces_h Number of faces in the height dimension
+ * @tparam face_h The height index of the face
+ * @tparam face_w The width index of the face
+ * @param tile_ptr Pointer to the start of the tile
+ * @param fill_value The value to use for padding
+ */
+template <
+    typename T,
+    uint32_t num_elements_unpadded_w,
+    uint32_t num_elements_unpadded_h,
+    uint32_t num_faces_w,
+    uint32_t num_faces_h,
+    uint32_t face_h,
+    uint32_t face_w>
+void fill_pad_face(T* tile_ptr, uint32_t fill_value) {
+    // Calculate face offset
+    constexpr uint32_t face_offset = (face_h * num_faces_w + face_w) * FACE_HW;
+    auto face_ptr = tile_ptr + face_offset;
+
+    // Right padding (width padding)
+    constexpr uint32_t face_w_offset = face_w * FACE_WIDTH;
+    constexpr uint32_t face_pad_w = (num_elements_unpadded_w <= face_w_offset) ? FACE_WIDTH
+                                    : (num_elements_unpadded_w >= face_w_offset + FACE_WIDTH)
+                                        ? 0
+                                        : face_w_offset + FACE_WIDTH - num_elements_unpadded_w;
+
+    if constexpr (face_pad_w > 0) {
+#pragma unroll
+        for (uint32_t row = 0; row < FACE_HEIGHT; ++row) {
+            auto row_ptr = face_ptr + row * FACE_WIDTH;
+            for (uint32_t col = FACE_WIDTH - face_pad_w; col < FACE_WIDTH; ++col) {
+                row_ptr[col] = fill_value;
+            }
+        }
+    }
+
+    // Bottom padding (height padding)
+    constexpr uint32_t face_h_offset = face_h * FACE_HEIGHT;
+    constexpr uint32_t face_pad_h = (num_elements_unpadded_h <= face_h_offset) ? FACE_HEIGHT
+                                    : (num_elements_unpadded_h >= face_h_offset + FACE_HEIGHT)
+                                        ? 0
+                                        : face_h_offset + FACE_HEIGHT - num_elements_unpadded_h;
+
+    if constexpr (face_pad_h > 0) {
+#pragma unroll
+        for (uint32_t row = FACE_HEIGHT - face_pad_h; row < FACE_HEIGHT; ++row) {
+            auto row_ptr = face_ptr + row * FACE_WIDTH;
+            for (uint32_t col = 0; col < FACE_WIDTH; ++col) {
+                row_ptr[col] = fill_value;
+            }
+        }
+    }
+}
+
+/**
+ * @brief Fills padding regions in a tile with a specified value.
+ *
+ * This function processes a tile by dividing it into faces and applying padding to each face
+ * based on the unpadded dimensions. The padding is applied to both width and height dimensions
+ * as needed.
+ *
+ * The function uses template metaprogramming to unroll the face processing loops at compile time,
+ * making it efficient for hardware execution. It processes each face in the tile by:
+ * 1. Calculating the face offset in the tile
+ * 2. Determining the padding requirements for width and height
+ * 3. Applying the padding with the specified fill value
+ *
+ * @tparam T The type of the elements in the tile
+ * @tparam num_elements_unpadded_w Number of unpadded elements in the width dimension
+ * @tparam num_elements_unpadded_h Number of unpadded elements in the height dimension
+ * @tparam num_faces_w Number of faces in the width dimension (default: TILE_WIDTH / FACE_WIDTH)
+ * @tparam num_faces_h Number of faces in the height dimension (default: TILE_HEIGHT / FACE_HEIGHT)
+ * @param l1_tile_ptr Pointer to the start of the tile in L1 memory
+ * @param fill_value The value to use for padding
+ */
+template <
+    typename T,
+    uint32_t num_elements_unpadded_w,
+    uint32_t num_elements_unpadded_h,
+    uint32_t num_faces_w = TILE_WIDTH / FACE_WIDTH,
+    uint32_t num_faces_h = TILE_HEIGHT / FACE_HEIGHT>
+void fill_pad_tile(uint32_t l1_tile_ptr, uint32_t fill_value) {
+    auto tile_ptr = reinterpret_cast<T*>(l1_tile_ptr);
+
+    // Face 0, 0
+    fill_pad_face<T, num_elements_unpadded_w, num_elements_unpadded_h, num_faces_w, num_faces_h, 0, 0>(
+        tile_ptr, fill_value);
+
+    // Face 0, 1
+    if constexpr (num_faces_w > 1) {
+        fill_pad_face<T, num_elements_unpadded_w, num_elements_unpadded_h, num_faces_w, num_faces_h, 0, 1>(
+            tile_ptr, fill_value);
+    }
+
+    // Face 1, 0
+    if constexpr (num_faces_h > 1) {
+        fill_pad_face<T, num_elements_unpadded_w, num_elements_unpadded_h, num_faces_w, num_faces_h, 1, 0>(
+            tile_ptr, fill_value);
+    }
+
+    // Face 1, 1
+    if constexpr (num_faces_w > 1 && num_faces_h > 1) {
+        fill_pad_face<T, num_elements_unpadded_w, num_elements_unpadded_h, num_faces_w, num_faces_h, 1, 1>(
+            tile_ptr, fill_value);
+    }
+}
+
+/**
+ * @brief Pads the last K tile in a matrix multiplication operation.
+ *
+ * This function handles padding for the last K tile in a matrix multiplication operation.
+ * It applies zero padding based on the specified data format (Float32 or Float16_b) and
+ * the unpadded width of the last K tile.
+ *
+ * @tparam in0_data_format The data format of the input tensor (Float32 or Float16_b)
+ * @tparam in0_last_ktile_w The unpadded width of the last K tile
+ * @param l1_write_addr_in0 The L1 memory address where the zeros should be written
+ */
+template <DataFormat in0_data_format, uint32_t in0_last_ktile_w>
+void pad_last_ktile(uint32_t l1_write_addr_in0) {
+    if constexpr (in0_data_format == DataFormat::Float32) {
+        fill_pad_tile<uint32_t, in0_last_ktile_w, /*num_elements_unpadded_h=*/TILE_HEIGHT>(
+            l1_write_addr_in0, /*pad_value=*/0);
+    } else if constexpr (in0_data_format == DataFormat::Float16_b) {
+        fill_pad_tile<uint16_t, in0_last_ktile_w, /*num_elements_unpadded_h=*/TILE_HEIGHT>(
+            l1_write_addr_in0, /*pad_value=*/0);
+    }
+}
diff --git a/ttnn/cpp/ttnn/operations/matmul/device/kernels/dataflow/reader_bmm_8bank_output_tiles_partitioned.cpp b/ttnn/cpp/ttnn/operations/matmul/device/kernels/dataflow/reader_bmm_8bank_output_tiles_partitioned.cpp
index b18907e2652b..633d39feddfe 100644
--- a/ttnn/cpp/ttnn/operations/matmul/device/kernels/dataflow/reader_bmm_8bank_output_tiles_partitioned.cpp
+++ b/ttnn/cpp/ttnn/operations/matmul/device/kernels/dataflow/reader_bmm_8bank_output_tiles_partitioned.cpp
@@ -1,10 +1,11 @@
-// SPDX-FileCopyrightText: © 2023 Tenstorrent Inc.
+// SPDX-FileCopyrightText: © 2023 Tenstorrent AI ULC
 //
 // SPDX-License-Identifier: Apache-2.0
 
 #include <stdint.h>
 
 #include "dataflow_api.h"
+#include "pad_tile.hpp"
 
 void kernel_main() {
     // same arg indices as in reader_binary_diff_lenghts for compat
@@ -23,6 +24,7 @@ void kernel_main() {
 
     constexpr bool src0_is_dram = get_compile_time_arg_val(0) == 1;
     constexpr bool src1_is_dram = get_compile_time_arg_val(1) == 1;
+    constexpr uint32_t in0_last_ktile_w = get_compile_time_arg_val(2);
 
     // DPRINT << "Mt=" << Mt << " Kt=" << Kt << " Nt=" << Nt << " MtKt=" << MtKt << "KtNt=" << KtNt << ENDL();
     // DPRINT << "src0=" << src0_addr << " src1=" << src1_addr << ENDL();
@@ -60,6 +62,11 @@ void kernel_main() {
                 uint32_t l1_write_addr_in0 = get_write_ptr(cb_id_in0);
                 noc_async_read_tile(itileA, s0, l1_write_addr_in0);
                 noc_async_read_barrier();
+                if constexpr (in0_last_ktile_w > 0) {
+                    if (kt == Kt - 1) {
+                        pad_last_ktile<in0_data_format, in0_last_ktile_w>(l1_write_addr_in0);
+                    }
+                }
                 cb_push_back(cb_id_in0, onetile);
             }
 
diff --git a/ttnn/cpp/ttnn/operations/matmul/device/kernels/dataflow/reader_bmm_tile_layout_in0.cpp b/ttnn/cpp/ttnn/operations/matmul/device/kernels/dataflow/reader_bmm_tile_layout_in0.cpp
index 7b1904fd2ce7..74597557ca9c 100644
--- a/ttnn/cpp/ttnn/operations/matmul/device/kernels/dataflow/reader_bmm_tile_layout_in0.cpp
+++ b/ttnn/cpp/ttnn/operations/matmul/device/kernels/dataflow/reader_bmm_tile_layout_in0.cpp
@@ -1,10 +1,11 @@
-// SPDX-FileCopyrightText: © 2023 Tenstorrent Inc.
+// SPDX-FileCopyrightText: © 2023 Tenstorrent AI ULC
 //
 // SPDX-License-Identifier: Apache-2.0
 
 #include <stdint.h>
 
 #include "dataflow_api.h"
+#include "pad_tile.hpp"
 
 void kernel_main() {
     // in0/in1 common args
@@ -30,6 +31,7 @@ void kernel_main() {
     // COMPILE TIME ARGS
     // interleaved accessor args
     constexpr bool in0_is_dram = get_compile_time_arg_val(0) == 1;
+    constexpr uint32_t last_ktile_w = get_compile_time_arg_val(1);
 
     constexpr uint32_t cb_id_in0 = 0;
 
@@ -39,8 +41,8 @@ void kernel_main() {
     cb_push_back(cb_id_in0, in0_num_tiles);
 #else
 
-    const uint32_t in0_single_tile_size_bytes = get_tile_size(cb_id_in0);
-    const DataFormat in0_data_format = get_dataformat(cb_id_in0);
+    constexpr uint32_t in0_single_tile_size_bytes = get_tile_size(cb_id_in0);
+    constexpr DataFormat in0_data_format = get_dataformat(cb_id_in0);
     constexpr const uint32_t in0_tile_hw = get_tile_hw(cb_id_in0);
 
     uint32_t l1_write_addr_in0;
@@ -61,6 +63,14 @@ void kernel_main() {
                 for (uint32_t w = 0; w < in0_block_w; ++w) {
                     noc_async_read_tile(in0_tensor_tile_id, s0, l1_write_addr_in0);
 
+                    // Zero out padded regions for the very last tile
+                    if constexpr (last_ktile_w > 0) {
+                        if ((block == num_blocks - 1) && (w == in0_block_w - 1)) {
+                            noc_async_read_barrier();
+                            pad_last_ktile<in0_data_format, last_ktile_w>(l1_write_addr_in0);
+                        }
+                    }
+
                     l1_write_addr_in0 += in0_single_tile_size_bytes;
                     in0_tensor_tile_id += in0_tensor_stride_w;
                 }
diff --git a/ttnn/cpp/ttnn/operations/matmul/device/kernels/dataflow/reader_bmm_tile_layout_in0_sender_dram_sharded.cpp b/ttnn/cpp/ttnn/operations/matmul/device/kernels/dataflow/reader_bmm_tile_layout_in0_sender_dram_sharded.cpp
index 355054f07b4f..1601bc4f35a8 100644
--- a/ttnn/cpp/ttnn/operations/matmul/device/kernels/dataflow/reader_bmm_tile_layout_in0_sender_dram_sharded.cpp
+++ b/ttnn/cpp/ttnn/operations/matmul/device/kernels/dataflow/reader_bmm_tile_layout_in0_sender_dram_sharded.cpp
@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: © 2024 Tenstorrent Inc.
+// SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC
 //
 // SPDX-License-Identifier: Apache-2.0
 
@@ -6,28 +6,30 @@
 
 #include "dataflow_api.h"
 #include "hostdevcommon/common_values.hpp"
+#include "pad_tile.hpp"
 
 void kernel_main() {
     // COMPILE TIME ARGS
     // in0 block args
     constexpr uint32_t in0_block_num_tiles = get_compile_time_arg_val(0);
     constexpr uint32_t in0_block_size_bytes = get_compile_time_arg_val(1);
+    constexpr uint32_t in0_last_ktile_w = get_compile_time_arg_val(2);
     // in0 mcast args
-    uint32_t in0_mcast_sender_semaphore_addr = get_semaphore(get_compile_time_arg_val(2));
-    uint32_t in0_mcast_receiver_semaphore_addr = get_semaphore(get_compile_time_arg_val(3));
-    constexpr uint32_t in0_mcast_num_dests = get_compile_time_arg_val(4);
-    constexpr uint32_t in0_mcast_num_cores = get_compile_time_arg_val(5);
+    uint32_t in0_mcast_sender_semaphore_addr = get_semaphore(get_compile_time_arg_val(3));
+    uint32_t in0_mcast_receiver_semaphore_addr = get_semaphore(get_compile_time_arg_val(4));
+    constexpr uint32_t in0_mcast_num_dests = get_compile_time_arg_val(5);
+    constexpr uint32_t in0_mcast_num_cores = get_compile_time_arg_val(6);
     // block args
-    constexpr uint32_t num_blocks = get_compile_time_arg_val(6);
+    constexpr uint32_t num_blocks = get_compile_time_arg_val(7);
     // in0 mcast args
-    constexpr uint32_t in0_mcast_dest_noc_start_x = get_compile_time_arg_val(7);
-    constexpr uint32_t in0_mcast_dest_noc_start_y = get_compile_time_arg_val(8);
-    constexpr uint32_t in0_mcast_dest_noc_end_x = get_compile_time_arg_val(9);
-    constexpr uint32_t in0_mcast_dest_noc_end_y = get_compile_time_arg_val(10);
+    constexpr uint32_t in0_mcast_dest_noc_start_x = get_compile_time_arg_val(8);
+    constexpr uint32_t in0_mcast_dest_noc_start_y = get_compile_time_arg_val(9);
+    constexpr uint32_t in0_mcast_dest_noc_end_x = get_compile_time_arg_val(10);
+    constexpr uint32_t in0_mcast_dest_noc_end_y = get_compile_time_arg_val(11);
     // in0 semaphore always valid
-    uint32_t in0_mcast_sender_valid_semaphore = get_semaphore(get_compile_time_arg_val(11));
+    uint32_t in0_mcast_sender_valid_semaphore = get_semaphore(get_compile_time_arg_val(12));
 
-    constexpr uint32_t num_blocks_per_shard = get_compile_time_arg_val(12);
+    constexpr uint32_t num_blocks_per_shard = get_compile_time_arg_val(13);
     constexpr uint32_t num_storage_cores = num_blocks / num_blocks_per_shard;
 
     // RUNTIME ARGS
@@ -37,16 +39,18 @@ void kernel_main() {
         return;
     }
     const uint32_t sender_id = get_arg_val<uint32_t>(1);
-    tt_l1_ptr uint32_t* in0_mcast_sender_noc_x = (tt_l1_ptr uint32_t*)(get_arg_addr(2));
-    tt_l1_ptr uint32_t* in0_mcast_sender_noc_y = (tt_l1_ptr uint32_t*)(get_arg_addr(2 + num_storage_cores));
+    const bool is_last_ktile_padded = static_cast<bool>(get_arg_val<uint32_t>(2));
+
+    tt_l1_ptr uint32_t* in0_mcast_sender_noc_x = (tt_l1_ptr uint32_t*)(get_arg_addr(3));
+    tt_l1_ptr uint32_t* in0_mcast_sender_noc_y = (tt_l1_ptr uint32_t*)(get_arg_addr(3 + num_storage_cores));
 
     const uint32_t sender_block_id = sender_id * num_blocks_per_shard;
 
     constexpr uint32_t cb_id_in0 = 0;
     constexpr uint32_t cb_id_in2 = 2;  // Sharded cb
 
-    const uint32_t in0_single_tile_size_bytes = get_tile_size(cb_id_in0);
-    const DataFormat in0_data_format = get_dataformat(cb_id_in0);
+    constexpr uint32_t in0_single_tile_size_bytes = get_tile_size(cb_id_in0);
+    constexpr DataFormat in0_data_format = get_dataformat(cb_id_in0);
 
     uint32_t l1_write_addr_in0;
 
@@ -91,6 +95,15 @@ void kernel_main() {
             // Now we have the block in the CB address, we can mcast to dests!
             uint64_t in0_multicast_data_addr = in0_multicast_data_noc | in0_start_address;
 
+            // Zero out padded regions for the very last tile
+            if constexpr (in0_last_ktile_w > 0) {
+                DPRINT << "is_last_ktile_padded: " << (uint32_t)is_last_ktile_padded << ENDL();
+                if (is_last_ktile_padded && (i == num_blocks_per_shard - 1)) {
+                    auto in0_last_ktile_ptr = local_read_addr + in0_block_size_bytes - in0_single_tile_size_bytes;
+                    pad_last_ktile<in0_data_format, in0_last_ktile_w>(in0_last_ktile_ptr);
+                }
+            }
+
 #ifndef SKIP_MCAST
             // num_dests must not include source, since we are NOT really doing a local copy!
             noc_async_write_multicast(
@@ -121,6 +134,14 @@ void kernel_main() {
                 noc_semaphore_set(in0_mcast_sender_semaphore_addr_ptr, 0);
 
                 uint64_t in0_multicast_data_addr = in0_multicast_data_noc | in0_start_address;
+
+                // Zero out padded regions for the very last tile
+                if constexpr (in0_last_ktile_w > 0) {
+                    if (is_last_ktile_padded && (block == num_blocks - 1)) {
+                        auto in0_last_ktile_ptr = local_read_addr + in0_block_size_bytes - in0_single_tile_size_bytes;
+                        pad_last_ktile<in0_data_format, in0_last_ktile_w>(in0_last_ktile_ptr);
+                    }
+                }
 #ifndef SKIP_MCAST
                 noc_async_write_multicast_loopback_src(
                     local_read_addr, in0_multicast_data_addr, in0_block_size_bytes, in0_mcast_num_cores, true, true);
diff --git a/ttnn/cpp/ttnn/operations/matmul/device/kernels/dataflow/reader_bmm_tile_layout_in0_sender_padding.cpp b/ttnn/cpp/ttnn/operations/matmul/device/kernels/dataflow/reader_bmm_tile_layout_in0_sender_padding.cpp
index 21372676fcf8..569a995d36d7 100644
--- a/ttnn/cpp/ttnn/operations/matmul/device/kernels/dataflow/reader_bmm_tile_layout_in0_sender_padding.cpp
+++ b/ttnn/cpp/ttnn/operations/matmul/device/kernels/dataflow/reader_bmm_tile_layout_in0_sender_padding.cpp
@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: © 2023 Tenstorrent Inc.
+// SPDX-FileCopyrightText: © 2023 Tenstorrent AI ULC
 //
 // SPDX-License-Identifier: Apache-2.0
 
@@ -7,6 +7,7 @@
 #include "dataflow_api.h"
 #include "hostdevcommon/common_values.hpp"
 #include "cpp/ttnn/operations/ccl/kernel_common/worker_sync_utils.hpp"
+#include "pad_tile.hpp"
 
 void kernel_main() {
     uint32_t rt_args_idx = 0;
@@ -35,23 +36,25 @@ void kernel_main() {
     constexpr uint32_t in0_block_w = get_compile_time_arg_val(5);
     constexpr uint32_t in0_block_h = get_compile_time_arg_val(6);
     constexpr uint32_t in0_block_num_tiles = get_compile_time_arg_val(7);
-    constexpr bool extract_shard_sub_blocks = (bool)get_compile_time_arg_val(8);
-    constexpr uint32_t shard_width_in_tiles = get_compile_time_arg_val(9);
-    constexpr uint32_t shard_height_in_tiles = get_compile_time_arg_val(10);
+    constexpr uint32_t in0_last_ktile_w = get_compile_time_arg_val(8);
+
+    constexpr bool extract_shard_sub_blocks = (bool)get_compile_time_arg_val(9);
+    constexpr uint32_t shard_width_in_tiles = get_compile_time_arg_val(10);
+    constexpr uint32_t shard_height_in_tiles = get_compile_time_arg_val(11);
     // in0/in1 common args
-    constexpr uint32_t num_blocks_inner_dim = get_compile_time_arg_val(11);
-    constexpr uint32_t num_blocks_w_dim = get_compile_time_arg_val(12);
-    constexpr uint32_t num_blocks_h_dim = get_compile_time_arg_val(13);
+    constexpr uint32_t num_blocks_inner_dim = get_compile_time_arg_val(12);
+    constexpr uint32_t num_blocks_w_dim = get_compile_time_arg_val(13);
+    constexpr uint32_t num_blocks_h_dim = get_compile_time_arg_val(14);
     // in0 mcast args
-    uint32_t in0_mcast_sender_semaphore_addr = get_semaphore(get_compile_time_arg_val(14));
-    uint32_t in0_mcast_receiver_semaphore_addr = get_semaphore(get_compile_time_arg_val(15));
-    constexpr uint32_t in0_mcast_num_dests = get_compile_time_arg_val(16);
-    constexpr uint32_t in0_mcast_num_cores = get_compile_time_arg_val(17);
+    uint32_t in0_mcast_sender_semaphore_addr = get_semaphore(get_compile_time_arg_val(15));
+    uint32_t in0_mcast_receiver_semaphore_addr = get_semaphore(get_compile_time_arg_val(16));
+    constexpr uint32_t in0_mcast_num_dests = get_compile_time_arg_val(17);
+    constexpr uint32_t in0_mcast_num_cores = get_compile_time_arg_val(18);
     // batch args
-    constexpr uint32_t MtKt = get_compile_time_arg_val(18);  // if 0
-    constexpr uint32_t batch = get_compile_time_arg_val(19);
+    constexpr uint32_t MtKt = get_compile_time_arg_val(19);  // if 0
+    constexpr uint32_t batch = get_compile_time_arg_val(20);
 
-    constexpr bool fuse_op = (bool)get_compile_time_arg_val(20);
+    constexpr bool fuse_op = (bool)get_compile_time_arg_val(21);
 
     MatmulOpReceiver fused_op_receiver;
     if constexpr (fuse_op) {
@@ -149,6 +152,15 @@ void kernel_main() {
                             if (bh < num_blocks_h_dim - 1 || h < last_block_h) {
                                 noc_async_read_tile(in0_tensor_tile_id, s0, l1_write_addr_in0);
                             }
+
+                            // Zero out padded regions for the very last tile
+                            if constexpr (in0_last_ktile_w > 0) {
+                                if ((block == num_blocks_inner_dim - 1) && (w == in0_block_w - 1)) {
+                                    noc_async_read_barrier();
+                                    pad_last_ktile<in0_data_format, in0_last_ktile_w>(l1_write_addr_in0);
+                                }
+                            }
+
                             l1_write_addr_in0 += in0_single_tile_size_bytes;
                             in0_tensor_tile_id += in0_tensor_stride_w;
                         }
diff --git a/ttnn/cpp/ttnn/operations/matmul/device/kernels/dataflow/reader_bmm_tile_layout_in0_sender_receiver_padding_block_sharded.cpp b/ttnn/cpp/ttnn/operations/matmul/device/kernels/dataflow/reader_bmm_tile_layout_in0_sender_receiver_padding_block_sharded.cpp
index 65f73289d4a6..05522d517671 100644
--- a/ttnn/cpp/ttnn/operations/matmul/device/kernels/dataflow/reader_bmm_tile_layout_in0_sender_receiver_padding_block_sharded.cpp
+++ b/ttnn/cpp/ttnn/operations/matmul/device/kernels/dataflow/reader_bmm_tile_layout_in0_sender_receiver_padding_block_sharded.cpp
@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: © 2023 Tenstorrent Inc.
+// SPDX-FileCopyrightText: © 2023 Tenstorrent AI ULC
 //
 // SPDX-License-Identifier: Apache-2.0
 
@@ -7,6 +7,7 @@
 #include "dataflow_api.h"
 #include "hostdevcommon/common_values.hpp"
 #include "cpp/ttnn/operations/ccl/kernel_common/worker_sync_utils.hpp"
+#include "pad_tile.hpp"
 
 void kernel_main() {
     constexpr bool core_has_output_block_work = (bool)get_compile_time_arg_val(0);
@@ -14,25 +15,27 @@ void kernel_main() {
 
     constexpr uint32_t in0_block_num_tiles = get_compile_time_arg_val(2);
     constexpr uint32_t in0_block_size_bytes = get_compile_time_arg_val(3);
+    constexpr uint32_t in0_last_ktile_w = get_compile_time_arg_val(4);
+
     // in0/in1 common args
-    constexpr uint32_t num_blocks_inner_dim = get_compile_time_arg_val(4);
-    constexpr uint32_t num_blocks_w_dim = get_compile_time_arg_val(5);
-    constexpr uint32_t num_blocks_h_dim = get_compile_time_arg_val(6);
+    constexpr uint32_t num_blocks_inner_dim = get_compile_time_arg_val(5);
+    constexpr uint32_t num_blocks_w_dim = get_compile_time_arg_val(6);
+    constexpr uint32_t num_blocks_h_dim = get_compile_time_arg_val(7);
     // in0 mcast args
-    uint32_t in0_mcast_sender_semaphore_addr = get_semaphore(get_compile_time_arg_val(7));
-    uint32_t in0_mcast_receiver_semaphore_addr = get_semaphore(get_compile_time_arg_val(8));
-    constexpr uint32_t in0_mcast_num_dests = get_compile_time_arg_val(9);
-    constexpr uint32_t in0_mcast_num_cores = get_compile_time_arg_val(10);
-    constexpr uint32_t num_x = get_compile_time_arg_val(11);
-    constexpr uint32_t num_y = get_compile_time_arg_val(12);
-    constexpr bool transpose_mcast = (bool)get_compile_time_arg_val(13);
-    constexpr uint32_t shard_width_in_tiles = get_compile_time_arg_val(14);
-    constexpr uint32_t shard_height_in_tiles = get_compile_time_arg_val(15);
-    constexpr uint32_t in0_block_w = get_compile_time_arg_val(16);
-    constexpr uint32_t in0_block_h = get_compile_time_arg_val(17);
-
-    constexpr uint32_t batch = get_compile_time_arg_val(18);
-    constexpr bool fuse_op = (bool)get_compile_time_arg_val(19);
+    uint32_t in0_mcast_sender_semaphore_addr = get_semaphore(get_compile_time_arg_val(8));
+    uint32_t in0_mcast_receiver_semaphore_addr = get_semaphore(get_compile_time_arg_val(9));
+    constexpr uint32_t in0_mcast_num_dests = get_compile_time_arg_val(10);
+    constexpr uint32_t in0_mcast_num_cores = get_compile_time_arg_val(11);
+    constexpr uint32_t num_x = get_compile_time_arg_val(12);
+    constexpr uint32_t num_y = get_compile_time_arg_val(13);
+    constexpr bool transpose_mcast = (bool)get_compile_time_arg_val(14);
+    constexpr uint32_t shard_width_in_tiles = get_compile_time_arg_val(15);
+    constexpr uint32_t shard_height_in_tiles = get_compile_time_arg_val(16);
+    constexpr uint32_t in0_block_w = get_compile_time_arg_val(17);
+    constexpr uint32_t in0_block_h = get_compile_time_arg_val(18);
+
+    constexpr uint32_t batch = get_compile_time_arg_val(19);
+    constexpr bool fuse_op = (bool)get_compile_time_arg_val(20);
 
     uint32_t rt_args_idx = 0;
     const uint32_t sender_id = get_arg_val<uint32_t>(rt_args_idx++);
@@ -166,9 +169,24 @@ void kernel_main() {
                             in0_tensor_current_inner_dim_block_start_addr += shard_read_width;
 
                             noc_async_read_barrier();
+
+                            if constexpr (in0_last_ktile_w > 0) {
+                                if ((block == num_blocks_inner_dim - 1)) {
+                                    auto in0_last_ktile_w_ptr = l1_write_extract_shard_in0 - in0_single_tile_size_bytes;
+                                    pad_last_ktile<in0_data_format, in0_last_ktile_w>(in0_last_ktile_w_ptr);
+                                }
+                            }
                         } else {
                             in0_tensor_read_addr = in0_tensor_current_inner_dim_block_start_addr;
                             in0_tensor_current_inner_dim_block_start_addr += in0_block_size_bytes;
+
+                            if constexpr (in0_last_ktile_w > 0) {
+                                if ((block == num_blocks_inner_dim - 1)) {
+                                    auto in0_last_ktile_w_ptr =
+                                        in0_tensor_read_addr + in0_block_size_bytes - in0_single_tile_size_bytes;
+                                    pad_last_ktile<in0_data_format, in0_last_ktile_w>(in0_last_ktile_w_ptr);
+                                }
+                            }
                         }
 
                         // wait until all in0 mcast destinations have atomically incremented the in0 semaphore_addr
diff --git a/ttnn/cpp/ttnn/operations/matmul/device/matmul_op_multi_core_program_factory.cpp b/ttnn/cpp/ttnn/operations/matmul/device/matmul_op_multi_core_program_factory.cpp
index 2b0ad7770125..98f3fa687837 100644
--- a/ttnn/cpp/ttnn/operations/matmul/device/matmul_op_multi_core_program_factory.cpp
+++ b/ttnn/cpp/ttnn/operations/matmul/device/matmul_op_multi_core_program_factory.cpp
@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: © 2023 Tenstorrent Inc.
+// SPDX-FileCopyrightText: © 2023 Tenstorrent AI ULC
 //
 // SPDX-License-Identifier: Apache-2.0
 
@@ -92,7 +92,9 @@ tt::tt_metal::operation::ProgramWithCallbacks matmul_multi_core(
 
     bool src0_is_dram = src0_buffer->buffer_type() == tt_metal::BufferType::DRAM;
     bool src1_is_dram = src1_buffer->buffer_type() == tt_metal::BufferType::DRAM;
-    std::vector<uint32_t> reader_compile_time_args = {(uint32_t)src0_is_dram, (uint32_t)src1_is_dram};
+    uint32_t last_ktile_w = a.get_logical_shape()[-1] % TILE_WIDTH;
+    std::vector<uint32_t> reader_compile_time_args = {
+        (uint32_t)src0_is_dram, (uint32_t)src1_is_dram, (uint32_t)last_ktile_w};
 
     bool dst_is_dram = dst_buffer->buffer_type() == tt_metal::BufferType::DRAM;
     std::vector<uint32_t> writer_compile_time_args = {(std::uint32_t)output_cb_index, (std::uint32_t)dst_is_dram};
diff --git a/ttnn/cpp/ttnn/operations/matmul/device/matmul_op_multi_core_reuse_mcast_1d_program_factory.cpp b/ttnn/cpp/ttnn/operations/matmul/device/matmul_op_multi_core_reuse_mcast_1d_program_factory.cpp
index 8a48ba8ef83d..afa77582f05d 100644
--- a/ttnn/cpp/ttnn/operations/matmul/device/matmul_op_multi_core_reuse_mcast_1d_program_factory.cpp
+++ b/ttnn/cpp/ttnn/operations/matmul/device/matmul_op_multi_core_reuse_mcast_1d_program_factory.cpp
@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: © 2023 Tenstorrent Inc.
+// SPDX-FileCopyrightText: © 2023 Tenstorrent AI ULC
 //
 // SPDX-License-Identifier: Apache-2.0
 
@@ -277,6 +277,7 @@ tt::tt_metal::operation::ProgramWithCallbacks create_program_mcast_in0(
 
     uint32_t in0_num_subblocks = (out_block_h / out_subblock_h);
     uint32_t in0_block_num_tiles = out_subblock_h * in0_block_w * in0_num_subblocks;
+    uint32_t in0_last_ktile_w = a.get_logical_shape()[-1] % in0_tile.get_tile_shape()[1];
 
     std::vector<uint32_t> in0_sender_compile_time_args;
     if (in0_is_sharded) {
@@ -286,6 +287,8 @@ tt::tt_metal::operation::ProgramWithCallbacks create_program_mcast_in0(
 
             (std::uint32_t)in0_block_num_tiles,                         // in0_block_num_tiles
             (std::uint32_t)in0_block_num_tiles * in0_single_tile_size,  // in0_block_size_bytes
+            (std::uint32_t)in0_last_ktile_w,
+
             // in0/in1 common args
             (std::uint32_t)num_blocks,        // num_blocks
             (std::uint32_t)out_num_blocks_x,  // num_blocks_x
@@ -320,9 +323,10 @@ tt::tt_metal::operation::ProgramWithCallbacks create_program_mcast_in0(
             (std::uint32_t)in0_block_w,          // in0_block_w
             (std::uint32_t)in0_block_h,          // in0_block_h
             (std::uint32_t)in0_block_num_tiles,  // in0_block_num_tiles
-            (std::uint32_t)false,                // extract_shard_sub_blocks (not used for interleaved)
-            (std::uint32_t)0,                    // shard_width_in_tiles (not used for interleaved)
-            (std::uint32_t)0,                    // shard_height_in_tiles (not used for interleaved)
+            (std::uint32_t)in0_last_ktile_w,
+            (std::uint32_t)false,  // extract_shard_sub_blocks (not used for interleaved)
+            (std::uint32_t)0,      // shard_width_in_tiles (not used for interleaved)
+            (std::uint32_t)0,      // shard_height_in_tiles (not used for interleaved)
             // in0/in1 common args
             (std::uint32_t)num_blocks,        // num_blocks
             (std::uint32_t)out_num_blocks_x,  // num_blocks_x
@@ -964,6 +968,8 @@ tt::tt_metal::operation::ProgramWithCallbacks create_program_mcast_in0(
 }
 
 tt::tt_metal::operation::ProgramWithCallbacks create_program_mcast_in1(
+    tt::tt_metal::Program& program,
+    const tt::tt_metal::Tensor& a,
     tt_metal::IDevice* device,
     MathFidelity math_fidelity,
     bool fp32_dest_acc_en,
@@ -1001,8 +1007,6 @@ tt::tt_metal::operation::ProgramWithCallbacks create_program_mcast_in1(
     // currently only support transpose of the full tile
     bool in1_transpose_tile = in1_tile.get_transpose_of_faces() && in1_tile.get_transpose_within_face();
 
-    tt_metal::Program program{};
-
     bool fuse_op = false;
 
     uint32_t num_blocks = K / in0_block_w;
@@ -1041,6 +1045,8 @@ tt::tt_metal::operation::ProgramWithCallbacks create_program_mcast_in1(
     }
     uint32_t in0_CB_size = in0_CB_tiles * in0_single_tile_size;
 
+    uint32_t in0_last_ktile_w = a.get_logical_shape()[-1] % in0_tile.get_tile_shape()[1];
+
     bool extract_shard_sub_blocks = false;
     uint32_t in0_shard_height_in_tiles = 0;
     uint32_t in0_shard_width_in_tiles = 0;
@@ -1134,6 +1140,8 @@ tt::tt_metal::operation::ProgramWithCallbacks create_program_mcast_in1(
         (std::uint32_t)in0_block_w,                // in0_block_w
         (std::uint32_t)in0_block_h,                // in0_block_h
         (std::uint32_t)in0_block_w * in0_block_h,  // in0_block_num_tiles
+        (std::uint32_t)in0_last_ktile_w,
+
         (std::uint32_t)extract_shard_sub_blocks,
         (std::uint32_t)in0_shard_width_in_tiles,
         (std::uint32_t)in0_shard_height_in_tiles,
@@ -2567,6 +2575,8 @@ tt::tt_metal::operation::ProgramWithCallbacks matmul_multi_core_reuse_mcast_1d_o
             fused_op_signaler);
     } else {
         return reuse_mcast_1d_optimized_helpers::create_program_mcast_in1(
+            program,
+            a,
             device,
             math_fidelity,
             fp32_dest_acc_en,
diff --git a/ttnn/cpp/ttnn/operations/matmul/device/matmul_op_multi_core_reuse_mcast_2d_program_factory.cpp b/ttnn/cpp/ttnn/operations/matmul/device/matmul_op_multi_core_reuse_mcast_2d_program_factory.cpp
index c8745e4e4a86..716efd3e0ed0 100644
--- a/ttnn/cpp/ttnn/operations/matmul/device/matmul_op_multi_core_reuse_mcast_2d_program_factory.cpp
+++ b/ttnn/cpp/ttnn/operations/matmul/device/matmul_op_multi_core_reuse_mcast_2d_program_factory.cpp
@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: © 2023 Tenstorrent Inc.
+// SPDX-FileCopyrightText: © 2023 Tenstorrent AI ULC
 //
 // SPDX-License-Identifier: Apache-2.0
 
@@ -36,6 +36,7 @@ tt::tt_metal::operation::ProgramWithCallbacks create_program_mcast_in0_in1(
     uint32_t K,
     bool bcast_batch,
     uint32_t in0_block_w,
+    uint32_t in0_last_ktile_w,
     uint32_t out_subblock_h,
     uint32_t out_subblock_w,
     uint32_t out_block_h,
@@ -334,6 +335,8 @@ tt::tt_metal::operation::ProgramWithCallbacks create_program_mcast_in0_in1(
 
             (std::uint32_t)in0_block_num_tiles,                         // in0_block_num_tiles
             (std::uint32_t)in0_block_num_tiles * in0_single_tile_size,  // in0_block_size_bytes
+            (std::uint32_t)in0_last_ktile_w,
+
             // in0/in1 common args
             (std::uint32_t)num_blocks,  // num_blocks
             (std::uint32_t)out_num_blocks_x,
@@ -364,9 +367,11 @@ tt::tt_metal::operation::ProgramWithCallbacks create_program_mcast_in0_in1(
             (std::uint32_t)in0_block_w,      // in0_tensor_next_inner_dim_block_stride
             (std::uint32_t)K * in0_block_h,  // in0_tensor_next_h_dim_block_stride
             // in0 block args
-            (std::uint32_t)in0_block_w,                // in0_block_w
-            (std::uint32_t)in0_block_h,                // in0_block_h
-            (std::uint32_t)in0_block_num_tiles,        // in0_block_num_tiles
+            (std::uint32_t)in0_block_w,          // in0_block_w
+            (std::uint32_t)in0_block_h,          // in0_block_h
+            (std::uint32_t)in0_block_num_tiles,  // in0_block_num_tiles
+            (std::uint32_t)in0_last_ktile_w,
+
             (std::uint32_t)false,                      // extract_shard_sub_blocks (not used for interleaved)
             (std::uint32_t)in0_shard_width_in_tiles,   // shard_width_in_tiles (not used for interleaved)
             (std::uint32_t)in0_shard_height_in_tiles,  // shard_height_in_tiles (not used for interleaved)
@@ -1355,6 +1360,8 @@ tt::tt_metal::operation::ProgramWithCallbacks matmul_multi_core_reuse_mcast_2d_o
     tt::DataFormat in1_data_format = tt_metal::datatype_to_dataformat_converter(b.get_dtype());          // in1
     tt::DataFormat output_data_format = tt_metal::datatype_to_dataformat_converter(output.get_dtype());  // output
 
+    uint32_t in0_last_ktile_w = a.get_logical_shape()[-1] % in0_tile.get_tile_shape()[1];
+
     tt_metal::Buffer* bias_buffer = nullptr;
     tt::DataFormat bias_data_format = tt::DataFormat::Bfp8_b;  // bias; doesn't matter if bias=nullptr
     if (bias.has_value()) {
@@ -1449,6 +1456,7 @@ tt::tt_metal::operation::ProgramWithCallbacks matmul_multi_core_reuse_mcast_2d_o
         Kt,
         bcast_batch,
         in0_block_w,
+        in0_last_ktile_w,
         out_subblock_h,
         out_subblock_w,
         out_block_h,
diff --git a/ttnn/cpp/ttnn/operations/matmul/device/matmul_op_multi_core_reuse_mcast_dram_sharded_program_factory.cpp b/ttnn/cpp/ttnn/operations/matmul/device/matmul_op_multi_core_reuse_mcast_dram_sharded_program_factory.cpp
index 9fb52d945b13..d330e404cc47 100644
--- a/ttnn/cpp/ttnn/operations/matmul/device/matmul_op_multi_core_reuse_mcast_dram_sharded_program_factory.cpp
+++ b/ttnn/cpp/ttnn/operations/matmul/device/matmul_op_multi_core_reuse_mcast_dram_sharded_program_factory.cpp
@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: © 2023 Tenstorrent Inc.
+// SPDX-FileCopyrightText: © 2024 Tenstorrent AI ULC
 //
 // SPDX-License-Identifier: Apache-2.0
 
@@ -66,6 +66,7 @@ tt::tt_metal::operation::ProgramWithCallbacks create_program_dram_sharded(
     uint32_t N,
     uint32_t K,
     uint32_t in0_block_w,
+    uint32_t in0_last_ktile_w,
     uint32_t per_core_M,
     uint32_t per_core_N_storage,
     std::optional<UnaryWithParam> fused_activation,
@@ -303,6 +304,7 @@ tt::tt_metal::operation::ProgramWithCallbacks create_program_dram_sharded(
     std::vector<uint32_t> in0_sender_compile_time_args = {
         (std::uint32_t)in0_block_num_tiles,                         // in0_block_num_tiles
         (std::uint32_t)in0_block_num_tiles * in0_single_tile_size,  // in0_block_size_bytes
+        (std::uint32_t)in0_last_ktile_w,                            // in0_last_ktile_w
         // in0 mcast args
         (std::uint32_t)in0_mcast_sender_semaphore_id,
         (std::uint32_t)in0_mcast_receiver_semaphore_id,
@@ -598,6 +600,8 @@ tt::tt_metal::operation::ProgramWithCallbacks create_program_dram_sharded(
 
         mm_in0_sender_args.push_back((std::uint32_t)worker_core_type);
         mm_in0_sender_args.push_back((std::uint32_t)sender_id);
+        mm_in0_sender_args.push_back(
+            (std::uint32_t)((core == all_storage_cores_vec.back()) and (in0_last_ktile_w > 0)));
         mm_in0_sender_args.insert(
             mm_in0_sender_args.end(), in0_mcast_sender_noc_x.begin(), in0_mcast_sender_noc_x.end());
         mm_in0_sender_args.insert(
@@ -619,6 +623,7 @@ tt::tt_metal::operation::ProgramWithCallbacks create_program_dram_sharded(
         uint32_t worker_core_type = 3;
         mm_in0_receiver_args.push_back((std::uint32_t)worker_core_type);
         mm_in0_receiver_args.push_back((std::uint32_t)0);
+        mm_in0_receiver_args.push_back((std::uint32_t)0);  // in0_last_ktile_w
         mm_in0_receiver_args.insert(
             mm_in0_receiver_args.end(), in0_mcast_sender_noc_x.begin(), in0_mcast_sender_noc_x.end());
         mm_in0_receiver_args.insert(
@@ -963,6 +968,7 @@ tt::tt_metal::operation::ProgramWithCallbacks matmul_multi_core_reuse_dram_shard
     uint32_t Mt = get_batch_size(ashape) * ashape[-2] / in0_tile_shape[0];
     uint32_t Kt = ashape[-1] / in0_tile_shape[1];
     uint32_t Nt = bshape[-1] / in1_tile_shape[1];
+    uint32_t in0_last_ktile_w = a.get_logical_shape()[-1] % in0_tile_shape[1];
 
     TT_FATAL(Kt % in0_block_w == 0, "Error");
 
@@ -987,6 +993,7 @@ tt::tt_metal::operation::ProgramWithCallbacks matmul_multi_core_reuse_dram_shard
         Nt,
         Kt,
         in0_block_w,
+        in0_last_ktile_w,
         per_core_M,
         per_core_N,
         std::move(fused_activation),
diff --git a/ttnn/cpp/ttnn/operations/matmul/device/matmul_op_multi_core_reuse_optimized_program_factory.cpp b/ttnn/cpp/ttnn/operations/matmul/device/matmul_op_multi_core_reuse_optimized_program_factory.cpp
index 9f45edd232dc..5680fcc6db6c 100644
--- a/ttnn/cpp/ttnn/operations/matmul/device/matmul_op_multi_core_reuse_optimized_program_factory.cpp
+++ b/ttnn/cpp/ttnn/operations/matmul/device/matmul_op_multi_core_reuse_optimized_program_factory.cpp
@@ -1,4 +1,4 @@
-// SPDX-FileCopyrightText: © 2023 Tenstorrent Inc.
+// SPDX-FileCopyrightText: © 2023 Tenstorrent AI ULC
 //
 // SPDX-License-Identifier: Apache-2.0
 
@@ -31,6 +31,7 @@ tt::tt_metal::operation::ProgramWithCallbacks create_program(
     uint32_t K,
     bool bcast_batch,
     uint32_t in0_block_w,
+    uint32_t in0_last_ktile_w,
     uint32_t out_subblock_h,
     uint32_t out_subblock_w,
     uint32_t per_core_M,
@@ -163,6 +164,7 @@ tt::tt_metal::operation::ProgramWithCallbacks create_program(
     std::vector<uint32_t> reader_compile_time_args = {
         // interleaved accessor args
         (std::uint32_t)in0_is_dram,
+        (std::uint32_t)in0_last_ktile_w,
     };
     std::vector<uint32_t> reader_writer_compile_time_args = {// interleaved accessor args
                                                              (std::uint32_t)in1_is_dram,
@@ -532,6 +534,8 @@ tt::tt_metal::operation::ProgramWithCallbacks matmul_multi_core_reuse_optimized_
     uint32_t Kt = ashape[-1] / in0_tile_shape[1];
     uint32_t Nt = bshape[-1] / in1_tile_shape[1];
 
+    uint32_t in0_last_ktile_w = a.get_logical_shape()[-1] % in0_tile_shape[1];
+
     // TODO: Generalize
     TT_FATAL(!fuse_batch, "Only fuse_batch=false is supported for optimized bmm!");
 
@@ -566,6 +570,7 @@ tt::tt_metal::operation::ProgramWithCallbacks matmul_multi_core_reuse_optimized_
         Kt,
         bcast_batch,
         in0_block_w,
+        in0_last_ktile_w,
         out_subblock_h,
         out_subblock_w,
         per_core_M,


Reviews:
user: login: bbradelTT
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-21T18:45:14Z
user: login: nsorabaTT
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-21T19:21:58Z
user: login: nsorabaTT
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-21T19:24:10Z
user: login: nsorabaTT
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-21T19:46:13Z
user: login: nsorabaTT
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-21T20:37:16Z
user: login: yugaoTT
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-22T14:56:08Z
user: login: bbradelTT
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-22T15:07:24Z
user: login: nsorabaTT
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-22T15:45:13Z
user: login: nsorabaTT
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-22T15:46:21Z
user: login: nsorabaTT
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-23T07:36:23Z
user: login: nsorabaTT
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-23T07:36:27Z
user: login: nsorabaTT
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-23T07:36:33Z
user: login: nsorabaTT
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-23T07:36:59Z
user: login: nsorabaTT
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-23T07:37:49Z
user: login: bbradelTT
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-23T15:32:43Z
user: login: bbradelTT
body: 1. Please rebase and update the kernels that were put back in due to a revert.

Also, some of the kernels might be not used any more. You can ignore those or delete them.

2. You may also need to update some tests/models if the device and model pipelines show any changes.
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-23T15:34:52Z
user: login: yugaoTT
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-23T20:59:49Z
user: login: yugaoTT
state: CHANGES_REQUESTED
author_association: CONTRIBUTOR
submitted_at: 2025-05-23T21:01:44Z
user: login: yugaoTT
body: LGTM
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-05-28T14:11:22Z

Review comments:
diff_hunk: @@ -61,6 +63,11 @@ void kernel_main() {
                 for (uint32_t w = 0; w < in0_block_w; ++w) {
                     noc_async_read_tile(in0_tensor_tile_id, s0, l1_write_addr_in0);
 
+                    // Zero out padded regions for the very last tile
+                    if ((block == num_blocks - 1) && (w == in0_block_w - 1) && (last_ktile_w > 0)) {
path: ttnn/cpp/ttnn/operations/matmul/device/kernels/dataflow/reader_bmm_tile_layout_in0.cpp
user: login: bbradelTT
body: Can this be
```
if (constexpr lastk_ktile_w > 0) {
  if ((block == num_blocks - 1) && (w == in0_block_w - 1) {
```
and equivalent in other kernels.
Then I think the perf hit would be close to 0 if there is no padding.
created_at: 2025-05-21T18:18:31Z
author_association: CONTRIBUTOR
diff_hunk: @@ -2207,3 +2207,43 @@ def test_sharded_matmul_with_multiple_out_block_values(device, out_block_h, out_
     )
     output_tensor = ttnn.to_torch(output_tensor)
     assert_with_pcc(torch_output_tensor, output_tensor, pcc=pcc)
+
+
+@pytest.mark.parametrize(
+    "input_a_shape,input_b_shape,input_a_value,input_b_value,input_a_reshape,input_b_reshape",
+    [
+        ((12,), (1, 12, 4, 4), 1.0, 2.0, (12, 1, 1), (12, 1, 16)),
+        ((16,), (1, 16, 2, 2), 1.0, 3.0, (16, 1, 1), (16, 1, 4)),
+        ((8,), (1, 8, 3, 3), 2.0, 1.0, (8, 1, 1), (8, 1, 9)),
path: tests/ttnn/unit_tests/operations/matmul/test_matmul.py
user: login: bbradelTT
body: It would be good to verify that all the different kernels were used.

I think it'd be easiest if you just created the different program configs.

created_at: 2025-05-21T18:21:15Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,107 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#pragma once
+
+#include <tt-metalium/constants.hpp>
+
+using namespace tt::constants;
+
+/*
+ * @brief Zero pads a part of a face in a tile, dictated by the number of elements to pad on the right and bottom edges
+ * of the face.
+ *
+ * This function zero pads a part of a face in a tile. It handles both width and height padding
+ * by setting the specified number of elements to zero on the right and bottom edges of the face.
+ *
+ * @tparam T The type of the elements in the tile
+ * @param num_elements_padded_w Number of elements to pad on the right edge (width)
+ * @param num_elements_padded_h Number of elements to pad on the bottom edge (height)
+ * @param tile_ptr Pointer to the start of the face in the tile
+ */
+template <typename T>
+void fill_pad_face(uint32_t num_elements_padded_w, uint32_t num_elements_padded_h, T* tile_ptr, T fill_value) {
+    // Right padding (width padding)
+    if (num_elements_padded_w > 0) {
+        for (uint32_t row = 0; row < FACE_HEIGHT; ++row) {
+            auto row_ptr = tile_ptr + row * FACE_WIDTH + (FACE_WIDTH - num_elements_padded_w);
+            for (uint32_t col = 0; col < num_elements_padded_w; ++col) {
path: ttnn/cpp/ttnn/operations/matmul/device/kernels/dataflow/pad_tile.hpp
user: login: bbradelTT
body: This isn't really the column.

Would it be possible to do
```
            auto row_ptr = tile_ptr + row * FACE_WIDTH;
            for (uint32_t col = FACE_WIDTH - num_elements_padded_w; col < FACE_WIDTH; ++col) {
```
?
created_at: 2025-05-21T18:41:58Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,107 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#pragma once
+
+#include <tt-metalium/constants.hpp>
+
+using namespace tt::constants;
+
+/*
+ * @brief Zero pads a part of a face in a tile, dictated by the number of elements to pad on the right and bottom edges
+ * of the face.
+ *
+ * This function zero pads a part of a face in a tile. It handles both width and height padding
+ * by setting the specified number of elements to zero on the right and bottom edges of the face.
+ *
+ * @tparam T The type of the elements in the tile
+ * @param num_elements_padded_w Number of elements to pad on the right edge (width)
+ * @param num_elements_padded_h Number of elements to pad on the bottom edge (height)
+ * @param tile_ptr Pointer to the start of the face in the tile
+ */
+template <typename T>
+void fill_pad_face(uint32_t num_elements_padded_w, uint32_t num_elements_padded_h, T* tile_ptr, T fill_value) {
+    // Right padding (width padding)
+    if (num_elements_padded_w > 0) {
+        for (uint32_t row = 0; row < FACE_HEIGHT; ++row) {
+            auto row_ptr = tile_ptr + row * FACE_WIDTH + (FACE_WIDTH - num_elements_padded_w);
+            for (uint32_t col = 0; col < num_elements_padded_w; ++col) {
+                row_ptr[col] = fill_value;
+            }
+        }
+    }
+
+    // Bottom padding (height padding)
+    if (num_elements_padded_h > 0) {
+        for (uint32_t row = FACE_HEIGHT - num_elements_padded_h; row < FACE_HEIGHT; ++row) {
+            auto row_ptr = tile_ptr + row * FACE_WIDTH;
+            for (uint32_t col = 0; col < FACE_WIDTH; ++col) {
+                row_ptr[col] = fill_value;
+            }
+        }
+    }
+}
+
+/**
+ * @brief Zero pads a tile by handling padding for each face within the tile.
+ *
+ * This function handles zero padding for a tile by breaking it down into faces and applying
+ * padding to each face individually. It calculates the appropriate padding for each face
+ * based on the unpadded dimensions and applies the padding using fill_pad_face.
+ *
+ * @tparam T The type of the elements in the tile
+ * @tparam num_faces_w Number of faces in the width dimension (default: TILE_WIDTH / FACE_WIDTH)
+ * @tparam num_faces_h Number of faces in the height dimension (default: TILE_HEIGHT / FACE_HEIGHT)
+ * @param num_elements_unpadded_w Number of unpadded elements in the width dimension
+ * @param num_elements_unpadded_h Number of unpadded elements in the height dimension
+ * @param l1_tile_ptr Pointer to the start of the tile
+ */
+template <typename T, uint32_t num_faces_w = TILE_WIDTH / FACE_WIDTH, uint32_t num_faces_h = TILE_HEIGHT / FACE_HEIGHT>
+void fill_pad_tile(
+    uint32_t num_elements_unpadded_w, uint32_t num_elements_unpadded_h, uint32_t l1_tile_ptr, uint32_t fill_value) {
+    auto tile_ptr = reinterpret_cast<T*>(l1_tile_ptr);
+
+    for (uint32_t face_h = 0; face_h < num_faces_h; ++face_h) {
+        for (uint32_t face_w = 0; face_w < num_faces_w; ++face_w) {
+            // Calculate face offset
+            uint32_t face_offset = (face_h * num_faces_w + face_w) * FACE_HW;
+            auto face_ptr = tile_ptr + face_offset;
+            // Calculate padding for this specific face
+            if (num_elements_unpadded_w > 0) {
+                uint32_t face_pad_w = num_elements_unpadded_w < (face_w * FACE_WIDTH)
path: ttnn/cpp/ttnn/operations/matmul/device/kernels/dataflow/pad_tile.hpp
user: login: bbradelTT
body: What are the values for face_pad_w if num_elements_unpadded_w is
1. 15
2. 17
across the different faces?
created_at: 2025-05-21T18:44:59Z
author_association: CONTRIBUTOR
diff_hunk: @@ -61,6 +63,11 @@ void kernel_main() {
                 for (uint32_t w = 0; w < in0_block_w; ++w) {
                     noc_async_read_tile(in0_tensor_tile_id, s0, l1_write_addr_in0);
 
+                    // Zero out padded regions for the very last tile
+                    if ((block == num_blocks - 1) && (w == in0_block_w - 1) && (last_ktile_w > 0)) {
path: ttnn/cpp/ttnn/operations/matmul/device/kernels/dataflow/reader_bmm_tile_layout_in0.cpp
user: login: nsorabaTT
body: Agree. Will make the change. FWIW, I had checked on godbolt and the compiler was able to optimize this out when a `constexpr` is involved in a conditonal.

Example: Check the absence of branch instructions when constexpr is introduced: https://godbolt.org/z/K7v8ezvEr



created_at: 2025-05-21T19:21:58Z
author_association: CONTRIBUTOR
diff_hunk: @@ -2207,3 +2207,43 @@ def test_sharded_matmul_with_multiple_out_block_values(device, out_block_h, out_
     )
     output_tensor = ttnn.to_torch(output_tensor)
     assert_with_pcc(torch_output_tensor, output_tensor, pcc=pcc)
+
+
+@pytest.mark.parametrize(
+    "input_a_shape,input_b_shape,input_a_value,input_b_value,input_a_reshape,input_b_reshape",
+    [
+        ((12,), (1, 12, 4, 4), 1.0, 2.0, (12, 1, 1), (12, 1, 16)),
+        ((16,), (1, 16, 2, 2), 1.0, 3.0, (16, 1, 1), (16, 1, 4)),
+        ((8,), (1, 8, 3, 3), 2.0, 1.0, (8, 1, 1), (8, 1, 9)),
path: tests/ttnn/unit_tests/operations/matmul/test_matmul.py
user: login: nsorabaTT
body: Ack, will be in the next commit.
created_at: 2025-05-21T19:24:10Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,107 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#pragma once
+
+#include <tt-metalium/constants.hpp>
+
+using namespace tt::constants;
+
+/*
+ * @brief Zero pads a part of a face in a tile, dictated by the number of elements to pad on the right and bottom edges
+ * of the face.
+ *
+ * This function zero pads a part of a face in a tile. It handles both width and height padding
+ * by setting the specified number of elements to zero on the right and bottom edges of the face.
+ *
+ * @tparam T The type of the elements in the tile
+ * @param num_elements_padded_w Number of elements to pad on the right edge (width)
+ * @param num_elements_padded_h Number of elements to pad on the bottom edge (height)
+ * @param tile_ptr Pointer to the start of the face in the tile
+ */
+template <typename T>
+void fill_pad_face(uint32_t num_elements_padded_w, uint32_t num_elements_padded_h, T* tile_ptr, T fill_value) {
+    // Right padding (width padding)
+    if (num_elements_padded_w > 0) {
+        for (uint32_t row = 0; row < FACE_HEIGHT; ++row) {
+            auto row_ptr = tile_ptr + row * FACE_WIDTH + (FACE_WIDTH - num_elements_padded_w);
+            for (uint32_t col = 0; col < num_elements_padded_w; ++col) {
+                row_ptr[col] = fill_value;
+            }
+        }
+    }
+
+    // Bottom padding (height padding)
+    if (num_elements_padded_h > 0) {
+        for (uint32_t row = FACE_HEIGHT - num_elements_padded_h; row < FACE_HEIGHT; ++row) {
+            auto row_ptr = tile_ptr + row * FACE_WIDTH;
+            for (uint32_t col = 0; col < FACE_WIDTH; ++col) {
+                row_ptr[col] = fill_value;
+            }
+        }
+    }
+}
+
+/**
+ * @brief Zero pads a tile by handling padding for each face within the tile.
+ *
+ * This function handles zero padding for a tile by breaking it down into faces and applying
+ * padding to each face individually. It calculates the appropriate padding for each face
+ * based on the unpadded dimensions and applies the padding using fill_pad_face.
+ *
+ * @tparam T The type of the elements in the tile
+ * @tparam num_faces_w Number of faces in the width dimension (default: TILE_WIDTH / FACE_WIDTH)
+ * @tparam num_faces_h Number of faces in the height dimension (default: TILE_HEIGHT / FACE_HEIGHT)
+ * @param num_elements_unpadded_w Number of unpadded elements in the width dimension
+ * @param num_elements_unpadded_h Number of unpadded elements in the height dimension
+ * @param l1_tile_ptr Pointer to the start of the tile
+ */
+template <typename T, uint32_t num_faces_w = TILE_WIDTH / FACE_WIDTH, uint32_t num_faces_h = TILE_HEIGHT / FACE_HEIGHT>
+void fill_pad_tile(
+    uint32_t num_elements_unpadded_w, uint32_t num_elements_unpadded_h, uint32_t l1_tile_ptr, uint32_t fill_value) {
+    auto tile_ptr = reinterpret_cast<T*>(l1_tile_ptr);
+
+    for (uint32_t face_h = 0; face_h < num_faces_h; ++face_h) {
+        for (uint32_t face_w = 0; face_w < num_faces_w; ++face_w) {
+            // Calculate face offset
+            uint32_t face_offset = (face_h * num_faces_w + face_w) * FACE_HW;
+            auto face_ptr = tile_ptr + face_offset;
+            // Calculate padding for this specific face
+            if (num_elements_unpadded_w > 0) {
+                uint32_t face_pad_w = num_elements_unpadded_w < (face_w * FACE_WIDTH)
path: ttnn/cpp/ttnn/operations/matmul/device/kernels/dataflow/pad_tile.hpp
user: login: nsorabaTT
body: Depending on the tiny tile params, there can be 1 or 2 faces in the tile along width dimension.

| num_faces_w | num_elements_unpadded_w | face_pad_w (face 0) | face_pad_w (face 1) |
|--------|--------|--------|--------|
| 2 | 15 | 1 | 16
| 2 | 17 | 0 | 15
| 1 | 15 | 1 | -
| 1 | 17 | - | - 
created_at: 2025-05-21T19:46:13Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,107 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#pragma once
+
+#include <tt-metalium/constants.hpp>
+
+using namespace tt::constants;
+
+/*
+ * @brief Zero pads a part of a face in a tile, dictated by the number of elements to pad on the right and bottom edges
+ * of the face.
+ *
+ * This function zero pads a part of a face in a tile. It handles both width and height padding
+ * by setting the specified number of elements to zero on the right and bottom edges of the face.
+ *
+ * @tparam T The type of the elements in the tile
+ * @param num_elements_padded_w Number of elements to pad on the right edge (width)
+ * @param num_elements_padded_h Number of elements to pad on the bottom edge (height)
+ * @param tile_ptr Pointer to the start of the face in the tile
+ */
+template <typename T>
+void fill_pad_face(uint32_t num_elements_padded_w, uint32_t num_elements_padded_h, T* tile_ptr, T fill_value) {
+    // Right padding (width padding)
+    if (num_elements_padded_w > 0) {
+        for (uint32_t row = 0; row < FACE_HEIGHT; ++row) {
+            auto row_ptr = tile_ptr + row * FACE_WIDTH + (FACE_WIDTH - num_elements_padded_w);
+            for (uint32_t col = 0; col < num_elements_padded_w; ++col) {
path: ttnn/cpp/ttnn/operations/matmul/device/kernels/dataflow/pad_tile.hpp
user: login: nsorabaTT
body: Ack, fix in next commit.
created_at: 2025-05-21T20:37:16Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,160 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#pragma once
+
+#include <tt-metalium/constants.hpp>
+
+using namespace tt::constants;
+
+/*
+ * @brief Zero pads a part of a face in a tile, dictated by the number of elements to pad on the right and bottom edges
+ * of the face.
+ *
+ * This function zero pads a part of a face in a tile. It handles both width and height padding
+ * by setting the specified number of elements to zero on the right and bottom edges of the face.
+ *
+ * @tparam T The type of the elements in the tile
+ * @param num_elements_padded_w Number of elements to pad on the right edge (width)
+ * @param num_elements_padded_h Number of elements to pad on the bottom edge (height)
+ * @param tile_ptr Pointer to the start of the face in the tile
+ */
+template <typename T>
+void fill_pad_face(uint32_t num_elements_padded_w, uint32_t num_elements_padded_h, T* tile_ptr, T fill_value) {
+    // Right padding (width padding)
+    if (num_elements_padded_w > 0) {
+        for (uint32_t row = 0; row < FACE_HEIGHT; ++row) {
+            auto row_ptr = tile_ptr + row * FACE_WIDTH;
+            for (uint32_t col = FACE_WIDTH - num_elements_padded_w; col < FACE_WIDTH; ++col) {
+                row_ptr[col] = fill_value;
+            }
+        }
+    }
+
+    // Bottom padding (height padding)
+    if (num_elements_padded_h > 0) {
+        for (uint32_t row = FACE_HEIGHT - num_elements_padded_h; row < FACE_HEIGHT; ++row) {
+            auto row_ptr = tile_ptr + row * FACE_WIDTH;
+            for (uint32_t col = 0; col < FACE_WIDTH; ++col) {
+                row_ptr[col] = fill_value;
+            }
+        }
+    }
+}
+
+/**
+ * @brief Zero pads a tile by handling padding for each face within the tile.
+ *
+ * This function handles zero padding for a tile by breaking it down into faces and applying
+ * padding to each face individually. It calculates the appropriate padding for each face
+ * based on the unpadded dimensions and applies the padding using fill_pad_face.
+
+    Case 1: All elements in face are padded (num_elements_unpadded_w <= face_w_offset)
+
+    num_elements_unpadded_w
+            v
+    +----------------+----------------+
+    | x x x          |                |
+    | x x x          |     Face       |
+    | x x x          |                |
+    +----------------+----------------+
+    ^                ^
+                face_w_offset
+
+    Case 2: All elements are unpadded (num_elements_unpadded_w >= face_w_offset + FACE_WIDTH)
+    +----------------+----------------+
+    | x x x x x x x x| x x            |
+    | x x Face  x x x| x x            |
+    | x x x x x x x x| x x            |
+    +----------------+----------------+
+    ^                ^
+    face_w_offset    face_w_offset + FACE_WIDTH
+
+    Case 3: Some elements are padded (face_w_offset < num_elements_unpadded_w < face_w_offset + FACE_WIDTH)
+            num_elements_unpadded_w
+                          v
+    +----------------+----------------+
+    | x x x x x x x x| x x            |
+    | x x x x x x x x| x x  Face      |
+    | x x x x x x x x| x x            |
+    +----------------+----------------+
+    ^                ^                ^
+                face_w_offset         face_w_offset + FACE_WIDTH
+
+ * @tparam T The type of the elements in the tile
+ * @tparam num_faces_w Number of faces in the width dimension (default: TILE_WIDTH / FACE_WIDTH)
+ * @tparam num_faces_h Number of faces in the height dimension (default: TILE_HEIGHT / FACE_HEIGHT)
+ * @param num_elements_unpadded_w Number of unpadded elements in the width dimension
+ * @param num_elements_unpadded_h Number of unpadded elements in the height dimension
+ * @param l1_tile_ptr Pointer to the start of the tile
+ */
+template <typename T, uint32_t num_faces_w = TILE_WIDTH / FACE_WIDTH, uint32_t num_faces_h = TILE_HEIGHT / FACE_HEIGHT>
+void fill_pad_tile(
+    uint32_t num_elements_unpadded_w, uint32_t num_elements_unpadded_h, uint32_t l1_tile_ptr, uint32_t fill_value) {
+    auto tile_ptr = reinterpret_cast<T*>(l1_tile_ptr);
+
+    for (uint32_t face_h = 0; face_h < num_faces_h; ++face_h) {
+        for (uint32_t face_w = 0; face_w < num_faces_w; ++face_w) {
+            // Calculate face offset
+            uint32_t face_offset = (face_h * num_faces_w + face_w) * FACE_HW;
+            auto face_ptr = tile_ptr + face_offset;
+            // Calculate padding for this specific face
+            if (num_elements_unpadded_w > 0) {
+                uint32_t face_w_offset = face_w * FACE_WIDTH;
+                uint32_t face_pad_w;
+
+                if (num_elements_unpadded_w <= face_w_offset) {
+                    // All elements in this face are padded ones
+                    face_pad_w = FACE_WIDTH;
+                } else if (num_elements_unpadded_w >= face_w_offset + FACE_WIDTH) {
+                    // All elements in this face are unpadded ones
+                    face_pad_w = 0;
+                } else {
+                    // Only some elements in this face are padded ones
+                    face_pad_w = face_w_offset + FACE_WIDTH - num_elements_unpadded_w;
+                }
+
+                fill_pad_face<T>(face_pad_w, 0, face_ptr, fill_value);
+            }
+
+            if (num_elements_unpadded_h > 0) {
+                uint32_t face_h_offset = face_h * FACE_HEIGHT;
+                uint32_t face_pad_h;
+
+                if (num_elements_unpadded_h <= face_h_offset) {
+                    // All elements in this face are padded ones
+                    face_pad_h = FACE_HEIGHT;
+                } else if (num_elements_unpadded_h >= face_h_offset + FACE_HEIGHT) {
+                    // All elements in this face are unpadded ones
+                    face_pad_h = 0;
+                } else {
+                    // Only some elements in this face are padded ones
+                    face_pad_h = face_h_offset + FACE_HEIGHT - num_elements_unpadded_h;
+                }
+                fill_pad_face<T>(0, face_pad_h, face_ptr, fill_value);
+            }
+        }
+    }
+}
+
+/**
+ * @brief Pads the last K tile in a matrix multiplication operation.
+ *
+ * This function handles padding for the last K tile in a matrix multiplication operation.
+ * It applies zero padding based on the specified data format (Float32 or Float16_b) and
+ * the unpadded width of the last K tile.
+ *
+ * @tparam in0_data_format The data format of the input tensor (Float32 or Float16_b)
+ * @param in0_last_ktile_w The unpadded width of the last K tile
+ * @param l1_write_addr_in0 The L1 memory address where the zeros should be written
+ */
+
+template <DataFormat in0_data_format>
+void pad_last_ktile(uint32_t in0_last_ktile_w, uint32_t l1_write_addr_in0) {
+    if constexpr (in0_data_format == DataFormat::Float32) {
+        fill_pad_tile<uint32_t>(in0_last_ktile_w, /*num_elements_unpadded_h=*/0, l1_write_addr_in0, /*pad_value=*/0);
+    } else if constexpr (in0_data_format == DataFormat::Float16_b) {
+        fill_pad_tile<uint16_t>(in0_last_ktile_w, /*num_elements_unpadded_h=*/0, l1_write_addr_in0, /*pad_value=*/0);
path: ttnn/cpp/ttnn/operations/matmul/device/kernels/dataflow/pad_tile.hpp
user: login: yugaoTT
body: can we just pad uint32 values instead? at least it will reduce cycles by 2
created_at: 2025-05-22T14:52:50Z
author_association: CONTRIBUTOR
diff_hunk: @@ -91,6 +94,12 @@ void kernel_main() {
             // Now we have the block in the CB address, we can mcast to dests!
             uint64_t in0_multicast_data_addr = in0_multicast_data_noc | in0_start_address;
 
+            // Zero out padded regions for the very last tile
+            if ((i == num_blocks_per_shard - 1) && (in0_last_ktile_w > 0)) {
path: ttnn/cpp/ttnn/operations/matmul/device/kernels/dataflow/reader_bmm_tile_layout_in0_sender_dram_sharded.cpp
user: login: yugaoTT
body: constexpr for in0_last_ktile_w so we don't regress the unpadded case. also not clear to me what does in0_last_ktile_w mean?
created_at: 2025-05-22T14:54:18Z
author_association: CONTRIBUTOR
diff_hunk: @@ -121,6 +130,12 @@ void kernel_main() {
                 noc_semaphore_set(in0_mcast_sender_semaphore_addr_ptr, 0);
 
                 uint64_t in0_multicast_data_addr = in0_multicast_data_noc | in0_start_address;
+
+                // Zero out padded regions for the very last tile
+                if ((block == num_blocks - 1) && (in0_last_ktile_w > 0)) {
path: ttnn/cpp/ttnn/operations/matmul/device/kernels/dataflow/reader_bmm_tile_layout_in0_sender_dram_sharded.cpp
user: login: yugaoTT
body: same above
created_at: 2025-05-22T14:54:38Z
author_association: CONTRIBUTOR
diff_hunk: @@ -149,6 +152,13 @@ void kernel_main() {
                             if (bh < num_blocks_h_dim - 1 || h < last_block_h) {
                                 noc_async_read_tile(in0_tensor_tile_id, s0, l1_write_addr_in0);
                             }
+
+                            // Zero out padded regions for the very last tile
+                            if ((block == num_blocks_inner_dim - 1) && (w == in0_block_w - 1) &&
path: ttnn/cpp/ttnn/operations/matmul/device/kernels/dataflow/reader_bmm_tile_layout_in0_sender_padding.cpp
user: login: yugaoTT
body: same above
created_at: 2025-05-22T14:54:49Z
author_association: CONTRIBUTOR
diff_hunk: @@ -166,6 +169,11 @@ void kernel_main() {
                             in0_tensor_current_inner_dim_block_start_addr += shard_read_width;
 
                             noc_async_read_barrier();
+
+                            if ((block == num_blocks_inner_dim - 1) && (in0_last_ktile_w > 0)) {
path: ttnn/cpp/ttnn/operations/matmul/device/kernels/dataflow/reader_bmm_tile_layout_in0_sender_receiver_padding_block_sharded.cpp
user: login: yugaoTT
body: same
created_at: 2025-05-22T14:55:27Z
author_association: CONTRIBUTOR
diff_hunk: @@ -277,6 +277,7 @@ tt::tt_metal::operation::ProgramWithCallbacks create_program_mcast_in0(
 
     uint32_t in0_num_subblocks = (out_block_h / out_subblock_h);
     uint32_t in0_block_num_tiles = out_subblock_h * in0_block_w * in0_num_subblocks;
+    uint32_t in0_last_ktile_w = a.get_logical_shape()[-1] % in0_tile.get_tile_shape()[1];
path: ttnn/cpp/ttnn/operations/matmul/device/matmul_op_multi_core_reuse_mcast_1d_program_factory.cpp
user: login: yugaoTT
body: is it just one tile actually need pad?
created_at: 2025-05-22T14:55:50Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,160 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#pragma once
+
+#include <tt-metalium/constants.hpp>
+
+using namespace tt::constants;
+
+/*
+ * @brief Zero pads a part of a face in a tile, dictated by the number of elements to pad on the right and bottom edges
+ * of the face.
+ *
+ * This function zero pads a part of a face in a tile. It handles both width and height padding
+ * by setting the specified number of elements to zero on the right and bottom edges of the face.
+ *
+ * @tparam T The type of the elements in the tile
+ * @param num_elements_padded_w Number of elements to pad on the right edge (width)
+ * @param num_elements_padded_h Number of elements to pad on the bottom edge (height)
+ * @param tile_ptr Pointer to the start of the face in the tile
+ */
+template <typename T>
+void fill_pad_face(uint32_t num_elements_padded_w, uint32_t num_elements_padded_h, T* tile_ptr, T fill_value) {
+    // Right padding (width padding)
+    if (num_elements_padded_w > 0) {
+        for (uint32_t row = 0; row < FACE_HEIGHT; ++row) {
+            auto row_ptr = tile_ptr + row * FACE_WIDTH;
+            for (uint32_t col = FACE_WIDTH - num_elements_padded_w; col < FACE_WIDTH; ++col) {
+                row_ptr[col] = fill_value;
+            }
+        }
+    }
+
+    // Bottom padding (height padding)
+    if (num_elements_padded_h > 0) {
+        for (uint32_t row = FACE_HEIGHT - num_elements_padded_h; row < FACE_HEIGHT; ++row) {
+            auto row_ptr = tile_ptr + row * FACE_WIDTH;
+            for (uint32_t col = 0; col < FACE_WIDTH; ++col) {
+                row_ptr[col] = fill_value;
+            }
+        }
+    }
+}
+
+/**
+ * @brief Zero pads a tile by handling padding for each face within the tile.
+ *
+ * This function handles zero padding for a tile by breaking it down into faces and applying
+ * padding to each face individually. It calculates the appropriate padding for each face
+ * based on the unpadded dimensions and applies the padding using fill_pad_face.
+
+    Case 1: All elements in face are padded (num_elements_unpadded_w <= face_w_offset)
+
+    num_elements_unpadded_w
+            v
+    +----------------+----------------+
+    | x x x          |                |
+    | x x x          |     Face       |
+    | x x x          |                |
+    +----------------+----------------+
+    ^                ^
+                face_w_offset
+
+    Case 2: All elements are unpadded (num_elements_unpadded_w >= face_w_offset + FACE_WIDTH)
+    +----------------+----------------+
+    | x x x x x x x x| x x            |
+    | x x Face  x x x| x x            |
+    | x x x x x x x x| x x            |
+    +----------------+----------------+
+    ^                ^
+    face_w_offset    face_w_offset + FACE_WIDTH
+
+    Case 3: Some elements are padded (face_w_offset < num_elements_unpadded_w < face_w_offset + FACE_WIDTH)
+            num_elements_unpadded_w
+                          v
+    +----------------+----------------+
+    | x x x x x x x x| x x            |
+    | x x x x x x x x| x x  Face      |
+    | x x x x x x x x| x x            |
+    +----------------+----------------+
+    ^                ^                ^
+                face_w_offset         face_w_offset + FACE_WIDTH
+
+ * @tparam T The type of the elements in the tile
+ * @tparam num_faces_w Number of faces in the width dimension (default: TILE_WIDTH / FACE_WIDTH)
+ * @tparam num_faces_h Number of faces in the height dimension (default: TILE_HEIGHT / FACE_HEIGHT)
+ * @param num_elements_unpadded_w Number of unpadded elements in the width dimension
+ * @param num_elements_unpadded_h Number of unpadded elements in the height dimension
+ * @param l1_tile_ptr Pointer to the start of the tile
+ */
+template <typename T, uint32_t num_faces_w = TILE_WIDTH / FACE_WIDTH, uint32_t num_faces_h = TILE_HEIGHT / FACE_HEIGHT>
+void fill_pad_tile(
+    uint32_t num_elements_unpadded_w, uint32_t num_elements_unpadded_h, uint32_t l1_tile_ptr, uint32_t fill_value) {
+    auto tile_ptr = reinterpret_cast<T*>(l1_tile_ptr);
+
+    for (uint32_t face_h = 0; face_h < num_faces_h; ++face_h) {
+        for (uint32_t face_w = 0; face_w < num_faces_w; ++face_w) {
+            // Calculate face offset
+            uint32_t face_offset = (face_h * num_faces_w + face_w) * FACE_HW;
+            auto face_ptr = tile_ptr + face_offset;
+            // Calculate padding for this specific face
+            if (num_elements_unpadded_w > 0) {
+                uint32_t face_w_offset = face_w * FACE_WIDTH;
+                uint32_t face_pad_w;
+
+                if (num_elements_unpadded_w <= face_w_offset) {
+                    // All elements in this face are padded ones
+                    face_pad_w = FACE_WIDTH;
+                } else if (num_elements_unpadded_w >= face_w_offset + FACE_WIDTH) {
+                    // All elements in this face are unpadded ones
+                    face_pad_w = 0;
+                } else {
+                    // Only some elements in this face are padded ones
+                    face_pad_w = face_w_offset + FACE_WIDTH - num_elements_unpadded_w;
+                }
+
+                fill_pad_face<T>(face_pad_w, 0, face_ptr, fill_value);
+            }
+
+            if (num_elements_unpadded_h > 0) {
+                uint32_t face_h_offset = face_h * FACE_HEIGHT;
+                uint32_t face_pad_h;
+
+                if (num_elements_unpadded_h <= face_h_offset) {
+                    // All elements in this face are padded ones
+                    face_pad_h = FACE_HEIGHT;
+                } else if (num_elements_unpadded_h >= face_h_offset + FACE_HEIGHT) {
+                    // All elements in this face are unpadded ones
+                    face_pad_h = 0;
+                } else {
+                    // Only some elements in this face are padded ones
+                    face_pad_h = face_h_offset + FACE_HEIGHT - num_elements_unpadded_h;
+                }
+                fill_pad_face<T>(0, face_pad_h, face_ptr, fill_value);
+            }
+        }
+    }
+}
+
+/**
+ * @brief Pads the last K tile in a matrix multiplication operation.
+ *
+ * This function handles padding for the last K tile in a matrix multiplication operation.
+ * It applies zero padding based on the specified data format (Float32 or Float16_b) and
+ * the unpadded width of the last K tile.
+ *
+ * @tparam in0_data_format The data format of the input tensor (Float32 or Float16_b)
+ * @param in0_last_ktile_w The unpadded width of the last K tile
+ * @param l1_write_addr_in0 The L1 memory address where the zeros should be written
+ */
+
+template <DataFormat in0_data_format>
+void pad_last_ktile(uint32_t in0_last_ktile_w, uint32_t l1_write_addr_in0) {
+    if constexpr (in0_data_format == DataFormat::Float32) {
+        fill_pad_tile<uint32_t>(in0_last_ktile_w, /*num_elements_unpadded_h=*/0, l1_write_addr_in0, /*pad_value=*/0);
+    } else if constexpr (in0_data_format == DataFormat::Float16_b) {
+        fill_pad_tile<uint16_t>(in0_last_ktile_w, /*num_elements_unpadded_h=*/0, l1_write_addr_in0, /*pad_value=*/0);
path: ttnn/cpp/ttnn/operations/matmul/device/kernels/dataflow/pad_tile.hpp
user: login: bbradelTT
body: If the padding count is odd, the last one would be uint16_t, right?
created_at: 2025-05-22T15:07:24Z
author_association: CONTRIBUTOR
diff_hunk: @@ -91,6 +94,12 @@ void kernel_main() {
             // Now we have the block in the CB address, we can mcast to dests!
             uint64_t in0_multicast_data_addr = in0_multicast_data_noc | in0_start_address;
 
+            // Zero out padded regions for the very last tile
+            if ((i == num_blocks_per_shard - 1) && (in0_last_ktile_w > 0)) {
path: ttnn/cpp/ttnn/operations/matmul/device/kernels/dataflow/reader_bmm_tile_layout_in0_sender_dram_sharded.cpp
user: login: nsorabaTT
body: 1. For the width sharded case, the `in0_last_ktile` is only applied to the core having the shard with last tile in k dimension. Hence, this is a runtime arg.

Example:
| S0 | S1 | S2 | S3 |
|--------|--------|--------|--------|

Only S3 has the last tile and will apply padding. While all of them will run the same kernel, the `in0_last_ktile_w` is different. So this cannot be made a constexpr.

For all other readers (e.g. height sharding, no sharding etc.), `in0_ktile_w` is a compile time arg and the constexpr is added.


2. `ktile` means the tile along `k` dimension for the matmul. `last` here refers to the last tile along k dimension. `w` here refers to the width of the tile having actual data (excluding padding). Putting together, `in0_last_ktile_w` refers to the number of elements in the last k tile that are valid.
created_at: 2025-05-22T15:45:13Z
author_association: CONTRIBUTOR
diff_hunk: @@ -277,6 +277,7 @@ tt::tt_metal::operation::ProgramWithCallbacks create_program_mcast_in0(
 
     uint32_t in0_num_subblocks = (out_block_h / out_subblock_h);
     uint32_t in0_block_num_tiles = out_subblock_h * in0_block_w * in0_num_subblocks;
+    uint32_t in0_last_ktile_w = a.get_logical_shape()[-1] % in0_tile.get_tile_shape()[1];
path: ttnn/cpp/ttnn/operations/matmul/device/matmul_op_multi_core_reuse_mcast_1d_program_factory.cpp
user: login: nsorabaTT
body: Only the last tile along `k` dimension. There may be multiple of these, depending on `M`.
created_at: 2025-05-22T15:46:21Z
author_association: CONTRIBUTOR
diff_hunk: @@ -61,6 +63,11 @@ void kernel_main() {
                 for (uint32_t w = 0; w < in0_block_w; ++w) {
                     noc_async_read_tile(in0_tensor_tile_id, s0, l1_write_addr_in0);
 
+                    // Zero out padded regions for the very last tile
+                    if ((block == num_blocks - 1) && (w == in0_block_w - 1) && (last_ktile_w > 0)) {
path: ttnn/cpp/ttnn/operations/matmul/device/kernels/dataflow/reader_bmm_tile_layout_in0.cpp
user: login: nsorabaTT
body: Addressed.
created_at: 2025-05-23T07:36:23Z
author_association: CONTRIBUTOR
diff_hunk: @@ -2207,3 +2207,43 @@ def test_sharded_matmul_with_multiple_out_block_values(device, out_block_h, out_
     )
     output_tensor = ttnn.to_torch(output_tensor)
     assert_with_pcc(torch_output_tensor, output_tensor, pcc=pcc)
+
+
+@pytest.mark.parametrize(
+    "input_a_shape,input_b_shape,input_a_value,input_b_value,input_a_reshape,input_b_reshape",
+    [
+        ((12,), (1, 12, 4, 4), 1.0, 2.0, (12, 1, 1), (12, 1, 16)),
+        ((16,), (1, 16, 2, 2), 1.0, 3.0, (16, 1, 1), (16, 1, 4)),
+        ((8,), (1, 8, 3, 3), 2.0, 1.0, (8, 1, 1), (8, 1, 9)),
path: tests/ttnn/unit_tests/operations/matmul/test_matmul.py
user: login: nsorabaTT
body: Addressed.
created_at: 2025-05-23T07:36:27Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,107 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#pragma once
+
+#include <tt-metalium/constants.hpp>
+
+using namespace tt::constants;
+
+/*
+ * @brief Zero pads a part of a face in a tile, dictated by the number of elements to pad on the right and bottom edges
+ * of the face.
+ *
+ * This function zero pads a part of a face in a tile. It handles both width and height padding
+ * by setting the specified number of elements to zero on the right and bottom edges of the face.
+ *
+ * @tparam T The type of the elements in the tile
+ * @param num_elements_padded_w Number of elements to pad on the right edge (width)
+ * @param num_elements_padded_h Number of elements to pad on the bottom edge (height)
+ * @param tile_ptr Pointer to the start of the face in the tile
+ */
+template <typename T>
+void fill_pad_face(uint32_t num_elements_padded_w, uint32_t num_elements_padded_h, T* tile_ptr, T fill_value) {
+    // Right padding (width padding)
+    if (num_elements_padded_w > 0) {
+        for (uint32_t row = 0; row < FACE_HEIGHT; ++row) {
+            auto row_ptr = tile_ptr + row * FACE_WIDTH + (FACE_WIDTH - num_elements_padded_w);
+            for (uint32_t col = 0; col < num_elements_padded_w; ++col) {
path: ttnn/cpp/ttnn/operations/matmul/device/kernels/dataflow/pad_tile.hpp
user: login: nsorabaTT
body: Addressed.
created_at: 2025-05-23T07:36:33Z
author_association: CONTRIBUTOR
diff_hunk: @@ -121,6 +130,12 @@ void kernel_main() {
                 noc_semaphore_set(in0_mcast_sender_semaphore_addr_ptr, 0);
 
                 uint64_t in0_multicast_data_addr = in0_multicast_data_noc | in0_start_address;
+
+                // Zero out padded regions for the very last tile
+                if ((block == num_blocks - 1) && (in0_last_ktile_w > 0)) {
path: ttnn/cpp/ttnn/operations/matmul/device/kernels/dataflow/reader_bmm_tile_layout_in0_sender_dram_sharded.cpp
user: login: nsorabaTT
body: For width sharded case, the variable is different value for different cores, so it is a runtime argument.
created_at: 2025-05-23T07:36:59Z
author_association: CONTRIBUTOR
diff_hunk: @@ -166,6 +169,11 @@ void kernel_main() {
                             in0_tensor_current_inner_dim_block_start_addr += shard_read_width;
 
                             noc_async_read_barrier();
+
+                            if ((block == num_blocks_inner_dim - 1) && (in0_last_ktile_w > 0)) {
path: ttnn/cpp/ttnn/operations/matmul/device/kernels/dataflow/reader_bmm_tile_layout_in0_sender_receiver_padding_block_sharded.cpp
user: login: nsorabaTT
body: Addressed.
created_at: 2025-05-23T07:37:49Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,160 @@
+// SPDX-FileCopyrightText: © 2025 Tenstorrent Inc.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+#pragma once
+
+#include <tt-metalium/constants.hpp>
+
+using namespace tt::constants;
+
+/*
+ * @brief Zero pads a part of a face in a tile, dictated by the number of elements to pad on the right and bottom edges
+ * of the face.
+ *
+ * This function zero pads a part of a face in a tile. It handles both width and height padding
+ * by setting the specified number of elements to zero on the right and bottom edges of the face.
+ *
+ * @tparam T The type of the elements in the tile
+ * @param num_elements_padded_w Number of elements to pad on the right edge (width)
+ * @param num_elements_padded_h Number of elements to pad on the bottom edge (height)
+ * @param tile_ptr Pointer to the start of the face in the tile
+ */
+template <typename T>
+void fill_pad_face(uint32_t num_elements_padded_w, uint32_t num_elements_padded_h, T* tile_ptr, T fill_value) {
+    // Right padding (width padding)
+    if (num_elements_padded_w > 0) {
+        for (uint32_t row = 0; row < FACE_HEIGHT; ++row) {
+            auto row_ptr = tile_ptr + row * FACE_WIDTH;
+            for (uint32_t col = FACE_WIDTH - num_elements_padded_w; col < FACE_WIDTH; ++col) {
+                row_ptr[col] = fill_value;
+            }
+        }
+    }
+
+    // Bottom padding (height padding)
+    if (num_elements_padded_h > 0) {
+        for (uint32_t row = FACE_HEIGHT - num_elements_padded_h; row < FACE_HEIGHT; ++row) {
+            auto row_ptr = tile_ptr + row * FACE_WIDTH;
+            for (uint32_t col = 0; col < FACE_WIDTH; ++col) {
+                row_ptr[col] = fill_value;
+            }
+        }
+    }
+}
+
+/**
+ * @brief Zero pads a tile by handling padding for each face within the tile.
+ *
+ * This function handles zero padding for a tile by breaking it down into faces and applying
+ * padding to each face individually. It calculates the appropriate padding for each face
+ * based on the unpadded dimensions and applies the padding using fill_pad_face.
+
+    Case 1: All elements in face are padded (num_elements_unpadded_w <= face_w_offset)
+
+    num_elements_unpadded_w
+            v
+    +----------------+----------------+
+    | x x x          |                |
+    | x x x          |     Face       |
+    | x x x          |                |
+    +----------------+----------------+
+    ^                ^
+                face_w_offset
+
+    Case 2: All elements are unpadded (num_elements_unpadded_w >= face_w_offset + FACE_WIDTH)
+    +----------------+----------------+
+    | x x x x x x x x| x x            |
+    | x x Face  x x x| x x            |
+    | x x x x x x x x| x x            |
+    +----------------+----------------+
+    ^                ^
+    face_w_offset    face_w_offset + FACE_WIDTH
+
+    Case 3: Some elements are padded (face_w_offset < num_elements_unpadded_w < face_w_offset + FACE_WIDTH)
+            num_elements_unpadded_w
+                          v
+    +----------------+----------------+
+    | x x x x x x x x| x x            |
+    | x x x x x x x x| x x  Face      |
+    | x x x x x x x x| x x            |
+    +----------------+----------------+
+    ^                ^                ^
+                face_w_offset         face_w_offset + FACE_WIDTH
+
+ * @tparam T The type of the elements in the tile
+ * @tparam num_faces_w Number of faces in the width dimension (default: TILE_WIDTH / FACE_WIDTH)
+ * @tparam num_faces_h Number of faces in the height dimension (default: TILE_HEIGHT / FACE_HEIGHT)
+ * @param num_elements_unpadded_w Number of unpadded elements in the width dimension
+ * @param num_elements_unpadded_h Number of unpadded elements in the height dimension
+ * @param l1_tile_ptr Pointer to the start of the tile
+ */
+template <typename T, uint32_t num_faces_w = TILE_WIDTH / FACE_WIDTH, uint32_t num_faces_h = TILE_HEIGHT / FACE_HEIGHT>
+void fill_pad_tile(
+    uint32_t num_elements_unpadded_w, uint32_t num_elements_unpadded_h, uint32_t l1_tile_ptr, uint32_t fill_value) {
+    auto tile_ptr = reinterpret_cast<T*>(l1_tile_ptr);
+
+    for (uint32_t face_h = 0; face_h < num_faces_h; ++face_h) {
+        for (uint32_t face_w = 0; face_w < num_faces_w; ++face_w) {
+            // Calculate face offset
+            uint32_t face_offset = (face_h * num_faces_w + face_w) * FACE_HW;
+            auto face_ptr = tile_ptr + face_offset;
+            // Calculate padding for this specific face
+            if (num_elements_unpadded_w > 0) {
+                uint32_t face_w_offset = face_w * FACE_WIDTH;
+                uint32_t face_pad_w;
+
+                if (num_elements_unpadded_w <= face_w_offset) {
+                    // All elements in this face are padded ones
+                    face_pad_w = FACE_WIDTH;
+                } else if (num_elements_unpadded_w >= face_w_offset + FACE_WIDTH) {
+                    // All elements in this face are unpadded ones
+                    face_pad_w = 0;
+                } else {
+                    // Only some elements in this face are padded ones
+                    face_pad_w = face_w_offset + FACE_WIDTH - num_elements_unpadded_w;
+                }
+
+                fill_pad_face<T>(face_pad_w, 0, face_ptr, fill_value);
+            }
+
+            if (num_elements_unpadded_h > 0) {
+                uint32_t face_h_offset = face_h * FACE_HEIGHT;
+                uint32_t face_pad_h;
+
+                if (num_elements_unpadded_h <= face_h_offset) {
+                    // All elements in this face are padded ones
+                    face_pad_h = FACE_HEIGHT;
+                } else if (num_elements_unpadded_h >= face_h_offset + FACE_HEIGHT) {
+                    // All elements in this face are unpadded ones
+                    face_pad_h = 0;
+                } else {
+                    // Only some elements in this face are padded ones
+                    face_pad_h = face_h_offset + FACE_HEIGHT - num_elements_unpadded_h;
+                }
+                fill_pad_face<T>(0, face_pad_h, face_ptr, fill_value);
+            }
+        }
+    }
+}
+
+/**
+ * @brief Pads the last K tile in a matrix multiplication operation.
+ *
+ * This function handles padding for the last K tile in a matrix multiplication operation.
+ * It applies zero padding based on the specified data format (Float32 or Float16_b) and
+ * the unpadded width of the last K tile.
+ *
+ * @tparam in0_data_format The data format of the input tensor (Float32 or Float16_b)
+ * @param in0_last_ktile_w The unpadded width of the last K tile
+ * @param l1_write_addr_in0 The L1 memory address where the zeros should be written
+ */
+
+template <DataFormat in0_data_format>
+void pad_last_ktile(uint32_t in0_last_ktile_w, uint32_t l1_write_addr_in0) {
+    if constexpr (in0_data_format == DataFormat::Float32) {
+        fill_pad_tile<uint32_t>(in0_last_ktile_w, /*num_elements_unpadded_h=*/0, l1_write_addr_in0, /*pad_value=*/0);
+    } else if constexpr (in0_data_format == DataFormat::Float16_b) {
+        fill_pad_tile<uint16_t>(in0_last_ktile_w, /*num_elements_unpadded_h=*/0, l1_write_addr_in0, /*pad_value=*/0);
path: ttnn/cpp/ttnn/operations/matmul/device/kernels/dataflow/pad_tile.hpp
user: login: bbradelTT
body: Just wanted to note.

Given that the overhead for larger inputs is 0.09% which is within noise, it's better not to complicate the code.
created_at: 2025-05-23T15:32:43Z
author_association: CONTRIBUTOR
diff_hunk: @@ -121,6 +130,12 @@ void kernel_main() {
                 noc_semaphore_set(in0_mcast_sender_semaphore_addr_ptr, 0);
 
                 uint64_t in0_multicast_data_addr = in0_multicast_data_noc | in0_start_address;
+
+                // Zero out padded regions for the very last tile
+                if ((block == num_blocks - 1) && (in0_last_ktile_w > 0)) {
path: ttnn/cpp/ttnn/operations/matmul/device/kernels/dataflow/reader_bmm_tile_layout_in0_sender_dram_sharded.cpp
user: login: yugaoTT
body: but we can introduce a define or another ct arg, to say okay we need padding, then go into this condition?
created_at: 2025-05-23T20:59:48Z
author_association: CONTRIBUTOR
diff_hunk: @@ -91,6 +94,12 @@ void kernel_main() {
             // Now we have the block in the CB address, we can mcast to dests!
             uint64_t in0_multicast_data_addr = in0_multicast_data_noc | in0_start_address;
 
+            // Zero out padded regions for the very last tile
+            if ((in0_last_ktile_w > 0) && (i == num_blocks_per_shard - 1)) {
path: ttnn/cpp/ttnn/operations/matmul/device/kernels/dataflow/reader_bmm_tile_layout_in0_sender_dram_sharded.cpp
user: login: yugaoTT
body: can wrap into a define or ct arg
created_at: 2025-05-23T21:00:51Z
author_association: CONTRIBUTOR
diff_hunk: @@ -121,6 +130,12 @@ void kernel_main() {
                 noc_semaphore_set(in0_mcast_sender_semaphore_addr_ptr, 0);
 
                 uint64_t in0_multicast_data_addr = in0_multicast_data_noc | in0_start_address;
+
+                // Zero out padded regions for the very last tile
+                if ((in0_last_ktile_w > 0) && (block == num_blocks - 1)) {
path: ttnn/cpp/ttnn/operations/matmul/device/kernels/dataflow/reader_bmm_tile_layout_in0_sender_dram_sharded.cpp
user: login: yugaoTT
body: can wrap into a define or ct arg
created_at: 2025-05-23T21:00:59Z
author_association: CONTRIBUTOR

PR comments:
user: login: bbradelTT
created_at: 2025-05-21T18:11:36Z
body: @nsorabaTT 
1. could there be any interaction with special values where having 0 is not sufficient (e.g. NaN or +-inf)?
2. what is the performance impact if the sizes are already multiples of 32?
author_association: CONTRIBUTOR
user: login: bbradelTT
created_at: 2025-05-21T18:13:56Z
body: 3. what is the performance impact if the sizes need to be padded?
author_association: CONTRIBUTOR
user: login: nsorabaTT
created_at: 2025-05-21T19:14:05Z
body: > @nsorabaTT
> 
> 1. could there be any interaction with special values where having 0 is not sufficient (e.g. NaN or +-inf)?
> 2. what is the performance impact if the sizes are already multiples of 32?
> 3. what is the performance impact if the sizes need to be padded?

1. Ran this experiment, presence of NaN/Inf in the in1 causes NaN to propagate to the output. Thanks for pointing this out. Looks like we will need to pad in1 as well.
2. For configs where entire K is read by all cores: no performance impact. This is because the argument is sent as a compile time argument to the reader kernel. When sizes are a multiple of 32, padding is disabled and so the compiler will optimize out the padding code.
For configs where input0 is width sharded, there is a small overhead of evaluating one if condition in the reader kernel. This is because only the storage core that contains the last tile does the padding. Rest all do not apply any padding. This overhead is minimal enough that I do not see any change in kernel execution time with this PR (for cases where there is no padding added).

3. I considered two methods to pad:
a. Use NOC transactions to move from MEM_ZEROS location to the target location.
b. Use RISCV cores to write to byte address.

(b) promises to be more performant as it will write to an entire 32B/64B address space in one go. However, (a) is more versatile as we can accommodate any write size. In this PR, I have opted to go with (b) since the write size can range from as small as 2B (1 BF16 element) to 124B (31 FP32 elements).


### Performance
As part of performance push, we can look at using approach (a) when an entire face is to be padded out.

Measuring the performance impact in percentage is not a good idea as the impact is reduced with longer K dimension.

For example, let us say it takes _T microseconds_ to pad 1 element. When K=1, 31 elements are padded. Thus, the increase is `(X + 31T) / X`

When K=131071, the time increase is `(Y + T) /T`.

I will update this comment in a bit with an average number for T in absolute sense to gauge the impact of padding.
author_association: CONTRIBUTOR
user: login: nsorabaTT
created_at: 2025-05-22T01:42:16Z
body: Tabulating the behavior of matmul for different floats, we see that the behavior is the same with or without padding. So we can conclude two things:

1. Padding just `in0` is sufficient to maintain existing behavior.
2. There needs to be a separate effort to have matmul family of ops match torch for NaN and inf inputs.


The below tables records whether the ttnn output matches torch. When it does not, the value in ttnn is provided with the torch value in braces.

| in0 | in1 | out (inputs unpadded) | output (inputs padded) |
|--------|--------|--------|--------|
| finite | finite | ✅ | ✅ |
| finite | zero | ✅  | ✅ |
| finite | inf | ✅ | ✅ |
| finite | NaN | inf (NaN) | inf (NaN) |
| zero | zero | ✅ | ✅ |
| zero | inf | 0 (NaN) | 0 (NaN) |
| zero | NaN | 0 (NaN) | 0 (NaN) |
| inf | inf | 0 (inf) | 0 (inf) |
| inf | NaN | 0 (NaN) | 0 (NaN) |
| NaN | NaN | 0 (NaN) | 0 (NaN) | 
author_association: CONTRIBUTOR
user: login: bbradelTT
created_at: 2025-05-22T14:48:15Z
body: @nsorabaTT in regards to whether padding just in0 is sufficient, we need to look at cases where in1 has values that are not padded. That can most easily be replicated by padding with those values.

E.g.
Case 1:
in0: logical values: finite, padded values: 0
in1: logical values: finite, garbage values (simulated by padding in1): inf


Case 2:
in0: logical values: finite, padded values: 0
in1: logical values: finite, garbage values (simulated by padding in1): NaN

author_association: CONTRIBUTOR
user: login: nsorabaTT
created_at: 2025-05-23T04:18:37Z
body: > @nsorabaTT in regards to whether padding just in0 is sufficient, we need to look at cases where in1 has values that are not padded. That can most easily be replicated by padding with those values.
> 
> E.g. Case 1: in0: logical values: finite, padded values: 0 in1: logical values: finite, garbage values (simulated by padding in1): inf
> 
> Case 2: in0: logical values: finite, padded values: 0 in1: logical values: finite, garbage values (simulated by padding in1): NaN

It has the same impact. Whether the NaN, +/-inf are present in the valid or padded region of in1, they contribute nothing to the dot product. I have explicitly run the sweep using `fill_implicit_padding` and here are the results:


| in0 | in1 | value in padded region of in1 | result correct? |
|--------|--------|--------|--------|
| finite | finite | finite | ✅ |
| finite | finite | zero | ✅ |
| finite | finite |  +/-inf | ✅ |
| finite | finite | NaN | ✅ |


Something that is concerning is the below outcome. Tracked separately: https://github.com/tenstorrent/tt-metal/issues/22517

| in0 | in1 | padded value of in1 | result correct? |
|--------|--------|--------|--------|
| finite | finite | Don't Care | ✅ |
| finite | infinite | Don't Care | ✅ |
| finite | NaN | Don't Care | ❌ |
author_association: CONTRIBUTOR
user: login: nsorabaTT
created_at: 2025-05-23T07:02:26Z
body: # Performance Analysis

We primarily look at two metrics:
1. What is the overhead introduced by the _extra code_ for existing cases where there is no padding
2. What is the increase in kernel duration due to padding being added to the last tile in K dimension.

## Code overhead
We analyze this by looking at kernel execution time with and without this PR. Tabulating it, we see its just noise.


| Config | % Increase* |
|--------|--------|
| MatmulMultiCoreReuseProgramConfig | 0.06% |
| MatmulMultiCoreReuseMultiCast1DProgramConfig | -0.04% | 
| MatmulMultiCoreReuseMultiCastDRAMShardedProgramConfig | 0.03% | 

* Increase in kernel duration for matmuls where there is no padding involved because of extra code introduced in this PR


## Padding overhead
We analyze this by looking at the incremental time added to kernel execution for each element that is padded. i.e. we measure kernel time for NCRISC by varying the number of valid elements in the last tile, ranging from 1 to 31.

### 1. Average time to pad an element: 170 ns
As expected, this number does not change with _K_ length. The time to pad an element in 170 ns when K=96 and K=9600.

### 2. What is the percentage increase?
It depends on K.

| K |% time increase per padded element |
|--------|--------|
| 96 | 9% |
| 9600 | 0.09% | 

Since the time remains independent of K, percentage increase is NOT a good measure.

### 3. Does this change depending on the MatMul Configuration used?
Not really. The variation is under 10%.

| Config | Time per padded element |
|--------|--------|
| MatmulMultiCoreReuseProgramConfig | 160 ns |
| MatmulMultiCoreReuseMultiCast1DProgramConfig | 178 ns | 
| MatmulMultiCoreReuseMultiCastDRAMShardedProgramConfig | 175 ns | 

### 5. Is the increase uniform?
Pretty much. Note the small drop when going from 31 to 32 elements. This is because for 32 elements, the padding code is also compiled away, reducing time.

![image](https://github.com/user-attachments/assets/ea584839-5ede-4762-bd70-9300c694d26d)


| Number of valid elements | NCRISC Kernel Duration (ns) |
|----------|----------------------------|
| 1 | 194314.20 |
| 2 | 194425.20 |
| 3 | 194114.20 |
| 4 | 194001.40 |
| 5 | 193821.10 |
| 6 | 193654.10 |
| 7 | 193470.30 |
| 8 | 193272.30 |
| 9 | 193041.40 |
| 10 | 193009.10 |
| 11 | 192930.10 |
| 12 | 192703.50 |
| 13 | 192525.40 |
| 14 | 192368.60 |
| 15 | 192381.40 |
| 16 | 192170.00 |
| 17 | 191866.20 |
| 18 | 191563.20 |
| 19 | 191543.20 |
| 20 | 191533.60 |
| 21 | 191357.70 |
| 22 | 191000.10 |
| 23 | 190891.60 |
| 24 | 190815.40 |
| 25 | 190750.10 |
| 26 | 190479.50 |
| 27 | 190535.10 |
| 28 | 190053.50 |
| 29 | 189937.30 |
| 30 | 189760.50 |
| 31 | 189890.10 |
| 32 | 189327.40 |


author_association: CONTRIBUTOR
user: login: yugaoTT
created_at: 2025-05-23T21:10:03Z
body: 175ns for padding 32 elements per tile, so 5.5ns per element (~5.5 cycles per store), this is a lot. Wonder we should unroll those padding loops, so that the stores can be pipelined
author_association: CONTRIBUTOR
user: login: nsorabaTT
created_at: 2025-05-23T21:15:28Z
body: > 175ns for padding 32 elements per tile, so 5.5ns per element (~5.5 cycles per store), this is a lot. Wonder we should unroll those padding loops, so that the stores can be pipelined

The stores have to be broken down into individual transactions because they won't align with 32B or 64B (for wormhole). So we cannot issue NOC transactions.

As for loop unrolling, we are looking at at variable padding. We will need to unroll for certain padding numbers and keep it as is for certain others. Did I get that correct?

Example: One when we pad an entire face (16x16), we can do some unrolling tricks to make the stores efficient. It will also be contiguous memory. When we pad 1-15 element in a face in all rows, I doubt there is much we can do.
author_association: CONTRIBUTOR
user: login: yugaoTT
created_at: 2025-05-23T21:23:51Z
body: oh I didn't mean noc txns, the riscv store is good for now. I think to unroll it, you just need all your passed in args to be either const or template? all those can be acquired from ct arg right?

say we pad 1-15 element in a face in all rows, 
can we do sth like
```
#pragma unroll
for (int I=1; I<16;i++) {
ptr[I] = 0;
}
```
author_association: CONTRIBUTOR
user: login: nsorabaTT
created_at: 2025-05-23T21:24:14Z
body: > 1. Please rebase and update the kernels that were put back in due to a revert.
> 
> Also, some of the kernels might be not used any more. You can ignore those or delete them.
> 
> 2. You may also need to update some tests/models if the device and model pipelines show any changes.

1 is complete.
2. No changes, device and model pipelines are clean.
author_association: CONTRIBUTOR
user: login: nsorabaTT
created_at: 2025-05-23T21:28:55Z
body: > oh I didn't mean noc txns, the riscv store is good for now. I think to unroll it, you just need all your passed in args to be either const or template? all those can be acquired from ct arg right?
> 
> say we pad 1-15 element in a face in all rows, can we do sth like
> 
> ```
> #pragma unroll
> for (int I=1; I<16;i++) {
> ptr[I] = 0;
> }
> ```

Good point. For the sharded case, I will separate out:
1. Whether padding is required (run time argument, bool)
2. If yes, how much padding is required (compile time argument)

Once this is done, padding is always a compile time argument so loop optimizations should kick in at compile time.
author_association: CONTRIBUTOR
user: login: nsorabaTT
created_at: 2025-05-27T23:09:22Z
body: > > oh I didn't mean noc txns, the riscv store is good for now. I think to unroll it, you just need all your passed in args to be either const or template? all those can be acquired from ct arg right?
> > say we pad 1-15 element in a face in all rows, can we do sth like
> > ```
> > #pragma unroll
> > for (int I=1; I<16;i++) {
> > ptr[I] = 0;
> > }
> > ```
> 
> Good point. For the sharded case, I will separate out:
> 
> 1. Whether padding is required (run time argument, bool)
> 2. If yes, how much padding is required (compile time argument)
> 
> Once this is done, padding is always a compile time argument so loop optimizations should kick in at compile time.

Resolved in the latest version.

@yugaoTT I will merge this, please let me know if you have more feedback, I will be happy to address it after I am back on 6/1.
author_association: CONTRIBUTOR