Skip to content

Comments

Fix errors in ttnn.nlp_create_qkv_heads_decode when input shard is bigger than one tile#19446

Merged
kpaigwar merged 2 commits intomainfrom
kp/createQKV_fix1
Mar 21, 2025
Merged

Fix errors in ttnn.nlp_create_qkv_heads_decode when input shard is bigger than one tile#19446
kpaigwar merged 2 commits intomainfrom
kp/createQKV_fix1

Conversation

@kpaigwar
Copy link
Contributor

@kpaigwar kpaigwar commented Mar 20, 2025

Problem description

Presently, CreateQKV Decode Op is only tested with hardcoded input shard shape (32, 32).
When relaxing this constraint and moving to bigger shard (32, 128). This results into two issues

  1. PCC issues when input is on contiguous core_grid
  2. FATAL error in creating output_specs when input is on subcoregrids.

What's changed

  1. Fixed bug in reader/writer kernel when calculating num_q_cores and num_kv_cores.
  2. Simplified create_output_specs, now fetching available output core_grid from output memory config.
  3. Automatic calculation of output core_grid, shard_shape, BufferType when no output_memory_config is specified.

Checklist

@kpaigwar kpaigwar force-pushed the kp/createQKV_fix1 branch 3 times, most recently from 05b1c1d to 80db6e0 Compare March 20, 2025 23:27
@kpaigwar kpaigwar marked this pull request as ready for review March 20, 2025 23:28
@kpaigwar kpaigwar force-pushed the kp/createQKV_fix1 branch from 80db6e0 to 3cc8414 Compare March 21, 2025 02:51
Copy link
Contributor

@cglagovichTT cglagovichTT left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

please update PR title and description with more detail

@kpaigwar kpaigwar changed the title fix createQKV for 10 cores Fix errors in ttnn.nlp_create_qkv_heads_decode when input shard is bigger than one tile Mar 21, 2025
@kpaigwar kpaigwar force-pushed the kp/createQKV_fix1 branch from 3cc8414 to b0f061a Compare March 21, 2025 14:11
@kpaigwar kpaigwar added this pull request to the merge queue Mar 21, 2025
Merged via the queue into main with commit ea87271 Mar 21, 2025
21 checks passed
@kpaigwar kpaigwar deleted the kp/createQKV_fix1 branch March 21, 2025 14:24
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants