Repository navigation
Wrong instruction dependency #69
Description
Describe the bug
In this nki kernel (https://gist.github.com/jinhongyii/492d7079efdf186b1cbb6b2f0c0478be), line 158,
p_transposed[p_loop_29[:, None, ], ((b_loop_2 * 128) + f_loop_22[None, :, ])] = nisa.tensor_copy(acc_psum_2[(block_kv * 126 + block_q) * 64 + b_loop_2, b_loop_2 % 8, p_loop_29[:, None, ], f_loop_22[None, :, ]])
this tensor_copy is accessing acc_psum_2 of bank b_loop_2 % 8, so the instruction in iteration b_loop_2 should have anti-dependency on iteration b_loop_2 - 8.
However, in the profile, it shows that this instruction depends on its previous iteration (b_loop_2 - 1). I believe this is a compiler bug which hurts performance.
Expected Behavior
The copy instruction should not depend on its previous iteration. The event semaphore in the screenshot should not exist.
Current Behavior
event semaphore appearing between each copy instruction under b_loop_2
Reproduction Steps
Run the nki kernel on trn1. You can reproduce the all the logs and profiles.
Regression Issue
- Select this option if this issue appears to be a regression.
Possible Solution
Fix the dependency building pass.
Additional Information/Context
No response
neuronx-cc version used
2.17.194
Framework(s) and their versions used (JAX, PyTorch, etc..)
No response
Activity
elaineyz commentedon Apr 22, 2025
Thank you for bringing this to our attention. We are investigating and will get back to you soon.
chongmni-aws commentedon Nov 1, 2025
Hi @jinhongyii, We've narrowed down the issues in the compiler scheduler and are actively working on the fix. Will get back to you once we have updates.
mrkcath-aws commentedon Oct 8, 2026
Re-tested on latest NKI on trn2, using a port of the b_loop_2 pattern (matmul into PSUM bank b % 8, then tensor_copy to SBUF). In the device profile, each tensor_copy waits only on the matmul that wrote its bank, and each matmul waits only on the copy from iteration b - 8 that last read the bank it overwrites. There is no wait between consecutive copies.
Closing as fixed; please reopen with a current repro if you still see it on the latest release.
github-actions commentedon Oct 8, 2026
This issue is now closed. Comments on closed issues are hard for our team to see.
If you need more assistance, please open a new issue that references this one.