Add support for unpadded shapes in Matmul1D w/ gather_in0 #16627
Conversation
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
avoraTT
requested review from
ayerofieiev-tt,
dmakoviichuk-tt,
rfurko-tt,
cfjchu,
TT-BrianLiu,
razorback3,
dongjin-na,
bbradelTT,
yugaoTT and
asandhupatlaTT
as code owners
yugaoTT
reviewed
Jan 10, 2025
ttnn/cpp/ttnn/operations/prefetcher/prefetcher/device/dram_prefetcher_op.cpp
Outdated
Show resolved
Hide resolved
yugaoTT
reviewed
Jan 10, 2025
tests/tt_eager/python_api_testing/unit_testing/misc/test_matmul_1d_gather_in0.py
Outdated
Show resolved
Hide resolved
avoraTT
force-pushed
the
avora/mm_pad
branch
3 times, most recently
from
414e20d to
a7aa7f3
Compare
bbradelTT
reviewed
Jan 13, 2025
ttnn/cpp/ttnn/operations/matmul/device/matmul_op_multi_core_reuse_mcast_1d_program_factory.cpp
Outdated
Show resolved
Hide resolved
bbradelTT
reviewed
Jan 14, 2025
avoraTT
requested review from
patrickroberts,
yan-zaretskiy,
eyonland,
abhullar-tt,
pgkeller,
aliuTT and
tt-aho
as code owners
avoraTT
force-pushed
the
avora/mm_pad
branch
2 times, most recently
from
bc32be0 to
1d004e5
Compare
yugaoTT
reviewed
Jan 21, 2025
...tnn/operations/matmul/device/kernels/dataflow/reader_bmm_tile_layout_in0_ring_all_gather.cpp
Outdated
Show resolved
Hide resolved
...erations/matmul/device/kernels/compute/bmm_large_block_zm_fused_bias_activation_gathered.cpp
Show resolved
Hide resolved
...tnn/operations/matmul/device/kernels/dataflow/reader_bmm_tile_layout_in0_ring_all_gather.cpp
Outdated
Show resolved
Hide resolved
ttnn/cpp/ttnn/operations/matmul/device/matmul_op_multi_core_reuse_mcast_1d_program_factory.cpp
Outdated
Show resolved
Hide resolved
avoraTT
force-pushed
the
avora/mm_pad
branch
2 times, most recently
from
72e3729 to
52628da
Compare
TT-BrianLiu
reviewed
Jan 21, 2025
bbradelTT
approved these changes
Jan 22, 2025
ttnn/cpp/ttnn/operations/prefetcher/prefetcher/device/dram_prefetcher_op.cpp
Outdated
Show resolved
Hide resolved
TT-BrianLiu
reviewed
Jan 22, 2025
Comment on lines
+1749
to
+1753
| /* Inner dim padding */ | ||
| const uint32_t Kt_pad = in0_buffer->shard_spec().shape()[1] / in0_tile.get_tile_shape()[1] * num_cores; | ||
| in0_block_w = Kt_pad / num_cores; | ||
|
|
||
| uint32_t num_blocks = Kt_pad / in0_block_w; |
There was a problem hiding this comment.
Should just use passed in in0_block_w and derive num_blocks with shard spec
There was a problem hiding this comment.
in0_block_w also needs to change to be derived from the padded value
There was a problem hiding this comment.
This code is just setting in0_block_w to be shard width. Is this not what you are passing in?
const uint32_t Kt_pad = in0_buffer->shard_spec().shape()[1] / in0_tile.get_tile_shape()[1] * num_cores;
in0_block_w = Kt_pad / num_cores;
num_blocks is basically unused in this function
There was a problem hiding this comment.
in0_block_w has to match with the unpadded K value, or else the other matmul validation fails. So that's why it needs to be updated here.
num_blocks is passed into the computer kernel to determine number of block iterations.
TT-BrianLiu
reviewed
Jan 22, 2025
ttnn/cpp/ttnn/operations/matmul/device/matmul_op_multi_core_reuse_mcast_1d_program_factory.cpp
Outdated
Show resolved
Hide resolved
avoraTT
force-pushed
the
avora/mm_pad
branch
2 times, most recently
from
f453ac9 to
3f13274
Compare
TT-BrianLiu
approved these changes
Jan 24, 2025
patrickroberts
pushed a commit
that referenced
this pull request
Jan 25, 2025
…6627) ### Ticket - #16626 ### Problem description In the current use case of Matmul1D with gather_in0 in the Llama models, the activations and weights need to be padded. This results in significant overhead. ### What's changed - Added support to skip part of in0_block_w that is padding information - Pad the Kt and Nt in the host code for gather_in0 ### Checklist - [x] Post commit CI passes (https://github.com/tenstorrent/tt-metal/actions/runs/12893880800) - [x] New/Existing tests provide coverage for changes (https://github.com/tenstorrent/tt-metal/actions/runs/12893883783)
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Ticket
Problem description
In the current use case of Matmul1D with gather_in0 in the Llama models, the activations and weights need to be padded. This results in significant overhead.
What's changed
Checklist