add tcgen.ld.red support to sm103a arch#2696
Conversation
stack-info: PR: #2696, branch: drisspg/stack/47
drisspg
changed the title
stack-info: PR: #2696, branch: drisspg/stack/47
drisspg
changed the title
stack-info: PR: #2696, branch: drisspg/stack/47
stack-info: PR: #2696, branch: drisspg/stack/47
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 9566d8d348
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| if const_expr(self.use_ldred_rowmax): | ||
| # ld.red returns each x32 tile's max in an extra register. | ||
| tSrS_red = cute.make_rmem_tensor(((1, 1), *tSrS_t2r.shape[1:]), self.qk_acc_dtype) | ||
| cute.copy(thr_tmem_load, tStS_t2r, (tSrS_t2r, tSrS_red)) |
There was a problem hiding this comment.
On SM103 dense kernels with more than one K tile, this ld.red load reads the same tStS_t2r TMEM address on each softmax_step while intervening MMA operations overwrite that slot. CUTLASS has lowered LdRed through side-effect-free inline asm unlike regular Ld32, so repeated reads from the same TMEM address can be CSE'd or hoisted and later tiles can reuse stale scores/max values, corrupting attention results. Please keep the normal TMEM load or use a side-effecting/fenced ld.red wrapper before enabling this path.
Useful? React with 👍 / 👎.
stack-info: PR: #2696, branch: drisspg/stack/47
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 1a9f9ebf86
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
drisspg
changed the title
drisspg
changed the title
drisspg
changed the title
1 participant
ad tcgen.ld.red support to sm103a arch
Inspired by; https://github.com/hao-ai-lab/flash-attention-fp4
Perf on dense non casual; makes sense msotly i dont really know why headim32 takes a hit though -> ill just update the expr to not dispatch for this
causal is similar