Support longer sequence lengths in ssm_prefix_scan
#9776
Merged
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
esmalTT
changed the title
ssm_prefix_scanssm_prefix_scan
esmalTT
force-pushed
the
esmal/prefix-scan-support-larger-seq
branch
4 times, most recently
from
6d4551f to
1f6f645
Compare
esmalTT
force-pushed
the
esmal/prefix-scan-support-larger-seq
branch
3 times, most recently
from
fc80603 to
d62da2c
Compare
kpaigwar
approved these changes
Jul 1, 2024
TT-BrianLiu
approved these changes
Jul 2, 2024
| @@ -4,10 +4,28 @@ | |||
|
|
|||
| #include "dataflow_api.h" | |||
|
|
|||
| void fill_zeros(uint32_t cb_id) { | |||
| constexpr uint32_t num_zeros_reads = 2048 / MEM_ZEROS_SIZE; | |||
There was a problem hiding this comment.
Should probably make 2048 a constexpr variable and comment that this is for bfloat16 only. Zeros CB for other data types will be different.
| @@ -88,7 +88,12 @@ operation::ProgramWithCallbacks multi_core_ssm_prefix_scan( | |||
| const uint32_t cb_zeros_id = tt::CB::c_intermed6; | |||
| const auto cb_zeros = create_circular_buffer(cb_zeros_id, 1, intermediary_tile_size, intermediary_format); | |||
|
|
|||
| std::vector<uint32_t> reader_compile_time_args = {cb_a_in_id, cb_bx_in_id}; | |||
| const uint32_t cb_h_acc_id = tt::CB::c_intermed7; | |||
| const uint32_t num_chunks_per_row = ceil(float(total_tiles_per_row) / 32.0f); | |||
There was a problem hiding this comment.
div_up might be cleaner for you without having to do these casts. Also use some variable for 32 const.
This change adds support for L > 32 in `ssm_eltwise_mul`. Logically this op can now handle any value of L but values of L > 128 will run out of L1 in `bfloat8` format. This change also fixes #9831.
esmalTT
force-pushed
the
esmal/prefix-scan-support-larger-seq
branch
from
d62da2c to
d5da023
Compare
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.
This PR adds support for L > 32 in
ssm_eltwise_mul. Logically we can handle any value of L but values of L > 128 will run out of L1 inbfloat8format.This change also addresses issue #9831.