[Schedule] Add an optional argument disable_checks for Schedule
#14281
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Motivation
Currently, some of the schedule checks are too strict, which makes it hard to schedule some workloads such as FlashAttention whose reduction is two-stage and does not strictly follows our standard.
This PR adds an optional argument
disable_checkswhich mutes some checks. The argument defaults toFalseand we can enable it whenever we want to disable somesoftchecks (bysoftwe mean if we violate such checks, the schedule is not necessarily invalid, and if we violatehardchecks the schedule step is invalid).In the future, we should collect the
softandhardchecks for all schedule primitives. This PR serves for FlashAttention and only caresbindand some reduction primitives for now.@spectrometerHBH @junrushao @Hzfengsy