[CuTe, FA4] Preserve first-tile flag during scheduler reconstruction - #2705
Conversation
dongxiao92
changed the title
dongxiao92
changed the title
37e2062 to
b509a9e
Compare
dongxiao92
changed the title
|
This changes should not be merged before cutlass-dsl 4.7 release. I will update the version once the dsl package is public |
|
@tridao 4.7.0 is live. I have also updated the version. Could you help to review? Thanks |
| values = values[n_items:] | ||
| return SingleTileScheduler(*(tuple(obj_list)), loc=self._loc) | ||
| scheduler = SingleTileScheduler(*(tuple(obj_list)), loc=self._loc) | ||
| # Preserve the Python-only first-tile flag across CuTe reconstruction. |
There was a problem hiding this comment.
is it safe to say this was a bug in older versions? regardless of 4.7 and this just happend to work; since I imagine we will need to sync the quack upgrade can we just have this change in this or a separate PR and the 4.7 bump isolated?
There was a problem hiding this comment.
is it safe to say this was a bug in older versions? regardless of 4.7 and this just happend to work
Yes. It happened to work and got exposed with some DSL frontend changes in 4.7.
can we just have this change in this or a separate PR and the 4.7 bump isolated?
Sounds good to me. I can cancel the version update. Should I add any comments to the code changes then?
There was a problem hiding this comment.
yeh I think 1 small commetn e.g. # Note python-only attributes; mlir drops the attributes that have no mlir equivalent and are only valid on the python object, and then for the later call sties you just do see # Note python-only attributes
There was a problem hiding this comment.
Done. Please help to review again
…ao-AILab#2705) * [CuTe, FA4] Preserve first-tile flag during scheduler reconstruction * Pin nvidia-cutlass-dsl to 4.7.0 * Keep scheduler fix separate from DSL upgrade (cherry picked from commit 7a08d7a) Signed-off-by: StevenWang-CY <stevenwang0805@outlook.com>
2 participants
Summary
Preserve the Python-only
_is_first_blockflag when CuTe reconstructs FlashAttention single-tile schedulers from their MLIR values.Root cause
These schedulers omit
_is_first_blockfrom their MLIR representation, but their constructors initialize it toTrue.After
advance_to_next_work()marks the first tile consumed, reconstruction therefore resets the omitted flag and makes the first tile valid again. The kernel can repeat that tile and deadlock its producer/consumer pipeline.Change
Copy
_is_first_blockfrom the original object after reconstructing:SingleTileSchedulerSingleTileVarlenSchedulerSm100FmhaStaticTileSchedulerValidation
Main branch and cutlass-4.7.0 could pass with this change