Skip to content

[CuTe, FA4] Preserve first-tile flag during scheduler reconstruction #2705

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
drisspg merged 5 commits into from
Aug 6, 2026
Merged

[CuTe, FA4] Preserve first-tile flag during scheduler reconstruction #2705

Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
b509a9e
[CuTe, FA4] Preserve first-tile flag during scheduler reconstruction
dongxiao92 Jul 13, 2026
98d0b4e
Merge main into agent/preserve-cute-scheduler-state
dongxiao92 Jul 21, 2026
ebe294a
Merge remote-tracking branch 'origin/main' into agent/preserve-cute-s…
dongxiao92 Aug 5, 2026
ac9ead4
Pin nvidia-cutlass-dsl to 4.7.0
dongxiao92 Aug 5, 2026
a7dc229
Keep scheduler fix separate from DSL upgrade
dongxiao92 Aug 6, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 13 additions & 3 deletions flash_attn/cute/tile_scheduler.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -342,7 +342,11 @@ def __new_from_mlir_values__(self, values):
for obj, n_items in zip([self.params, self._blk_coord], self._values_pos):
obj_list.append(cutlass.new_from_mlir_values(obj, values[:n_items]))
values = values[n_items:]
return SingleTileScheduler(*(tuple(obj_list)), loc=self._loc)
scheduler = SingleTileScheduler(*(tuple(obj_list)), loc=self._loc)
# Note: _is_first_block is a Python-only attribute omitted from MLIR values,
# so it must be restored explicitly after reconstruction.
scheduler._is_first_block = self._is_first_block
return scheduler


class StaticPersistentTileScheduler:
Expand Down Expand Up @@ -1385,7 +1389,10 @@ def __new_from_mlir_values__(self, values):
for obj, n_items in zip(objs, self._values_pos):
obj_list.append(cutlass.new_from_mlir_values(obj, values[:n_items]))
values = values[n_items:]
return self.__class__(*obj_list, loc=self._loc)
scheduler = self.__class__(*obj_list, loc=self._loc)
# See the note on Python-only attributes in SingleTileScheduler.
scheduler._is_first_block = self._is_first_block
return scheduler


class DynamicPersistentVarlenScheduler:
Expand Down Expand Up @@ -1827,9 +1834,12 @@ def __new_from_mlir_values__(self, values):
)
new_blk_coord = new_from_mlir_values(self._blk_coord, values[4:7])
new_grid_shape = new_from_mlir_values(self._grid_shape, values[7:])
return Sm100FmhaStaticTileScheduler(
scheduler = Sm100FmhaStaticTileScheduler(
new_params, new_current_work_linear_idx, new_blk_coord, new_grid_shape
)
# See the note on Python-only attributes in SingleTileScheduler.
scheduler._is_first_block = self._is_first_block
return scheduler


def compute_sm100_fmha_grid(
Expand Down