Add attention example in tracing test #44
Conversation
This PR adds an FP8 attention kernel to the tracing test to validate that it traces correctly.
There was a problem hiding this comment.
This is a good demonstrator for how attention could look in TKW.
This needs more iteration and requires features we do not yet support like tuple unpacking.
I think we do not need to have an end-to-end flow as prerequisite to add this tracing test, but we should at least rebase to use the Python operators we have and have a rough plan for supporting the assignments here through expansion + scheduling.
Let's for now convert this to a draft PR
| def div(lhs: "Register", rhs: "Register") -> "Register": | ||
| ... |
There was a problem hiding this comment.
The predeclarations here do not have CustomOps associated with them
We should map the div here to the Python operator like here. The branch needs to be rebased for that
| # CHECK-NEXT: unknown: sub | ||
| # CHECK-NEXT: unknown: sub_1 | ||
| # CHECK-NEXT: unknown: mul | ||
| # CHECK-NEXT: unknown: add |
There was a problem hiding this comment.
These will become proper sub, mul and add after rebasing on main
| # CHECK-NEXT: unknown: getitem | ||
| # CHECK-NEXT: unknown: getitem_1 | ||
| # CHECK-NEXT: unknown: getitem_2 |
There was a problem hiding this comment.
These getitem nodes stem from the tuple unpacking here: final_sum, _, final_output = repeat
We should think about how we want to handle these. I think for now it might make sense to intercept tracing for getitem on a CustomOp and emit a getResult op instead.
At least that makes sense in this case where the results of the reduction are accessed. I don't think we want to support this unpacking for general nodes so we should add further checks on when this is valid.
Anyway, I like the syntax you have here. Feels Pythonic to me
This PR adds an FP8 attention kernel to the
tracing test to validate that it traces correctly.