Conversation
Summary:
`num_stages=0` has been removed from TLX (and the AMD path is updated to
match), so every remaining use needs to move to `num_stages=1`. This diff is a
mechanical sweep of fbsource for all the places a `0` could reach a Triton
kernel.
Three mechanisms are covered:
1. `triton.Config(..., num_stages=0)` -- the kernel-level autotune param. This
is the bulk of the change: the TLX idiom of passing `num_stages=0` to
suppress the compiler's auto-pipelining on top of manual TLX pipelining.
Covers `hammer/v2`, `hammer/v3`, `ads_mkl/ops/tlx`,
`simplicial_attention`, `tritonbench` fusionbench, and the TLX tutorials.
2. `tl.range(..., num_stages=0)` -- the loop-level `tt.num_stages` attribute,
same intent, different knob.
3. Autotune sweeps that generate a 0 -- e.g. `for num_stages in [0, 1, 2]` in
`fbr/flash/triton`, the HIP-only `[0, 1]` sweeps in the ragged HSTU
attention scripts, `openfold_triton`, and FlagGems. The `0` is dropped
rather than rewritten to a duplicate `1` (`[0, 1, 2]` -> `[1, 2]`, not
`[1, 1, 2]`).
TLX docs and agent skills that instruct readers/agents to emit `num_stages=0`
(`RecGenHITL` language profiles, `kperfagent` TLX prompt skills, `ace`
kernel_info) are updated too, so regenerated kernels do not reintroduce it.
Deliberately not touched:
- `not_allowed_values = {0}` guards in `ads_mkl/.../hardware.py` and the `gem`
kernel library -- these already strip `0` from AMD sweeps and stay correct.
- TileLang kernels (`tilelang.jit` / `T.Kernel`) -- `num_stages` there is a
different DSL's parameter, not Triton's.
- `torch/_inductor/select_algorithm.py` -- a dummy sentinel object, never
launched.
- `llama4x`/`mslk` `triton_splitk.py` -- already at `num_stages = 1`; only a
stale `TODO` comment mentions `0`.
- `fbcode/gem/next_gen` -- an ACL'd path, so it lives in the child diff.
Phabricator rejects diffs that touch both ACL'd and non-ACL'd paths
(https://fburl.com/no-mixed-paths).
- `third-party/triton/stable` and `fbcode/triton_mtia/third_party/triton`
-- reverted per request; these vendored TLX tutorial copies should be
updated via the upstream sync instead.
- Generated snapshots under `pyper_models/*/archive/` and
`minimal_viable_ai/p4p/*/cloned_files/`, and unrelated `num_stages` concepts
(dataswarm, shardmanager, eval pipelines).
Note: D114743914 (`[Triton] [Addmm] Fix num_stages range`) may overlap; expect
to prune anything already covered by other diffs.
Reviewed By: jianyuh
Differential Revision: D114751669
Contributor
|
@njriasan has exported this pull request. If you are a Meta employee, you can view the originating Diff in D114751669. |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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.
Summary:
num_stages=0has been removed from TLX (and the AMD path is updated tomatch), so every remaining use needs to move to
num_stages=1. This diff is amechanical sweep of fbsource for all the places a
0could reach a Tritonkernel.
Three mechanisms are covered:
triton.Config(..., num_stages=0)-- the kernel-level autotune param. Thisis the bulk of the change: the TLX idiom of passing
num_stages=0tosuppress the compiler's auto-pipelining on top of manual TLX pipelining.
Covers
hammer/v2,hammer/v3,ads_mkl/ops/tlx,simplicial_attention,tritonbenchfusionbench, and the TLX tutorials.tl.range(..., num_stages=0)-- the loop-leveltt.num_stagesattribute,same intent, different knob.
Autotune sweeps that generate a 0 -- e.g.
for num_stages in [0, 1, 2]infbr/flash/triton, the HIP-only[0, 1]sweeps in the ragged HSTUattention scripts,
openfold_triton, and FlagGems. The0is droppedrather than rewritten to a duplicate
1([0, 1, 2]->[1, 2], not[1, 1, 2]).TLX docs and agent skills that instruct readers/agents to emit
num_stages=0(
RecGenHITLlanguage profiles,kperfagentTLX prompt skills,acekernel_info) are updated too, so regenerated kernels do not reintroduce it.
Deliberately not touched:
not_allowed_values = {0}guards inads_mkl/.../hardware.pyand thegemkernel library -- these already strip
0from AMD sweeps and stay correct.tilelang.jit/T.Kernel) --num_stagesthere is adifferent DSL's parameter, not Triton's.
torch/_inductor/select_algorithm.py-- a dummy sentinel object, neverlaunched.
llama4x/mslktriton_splitk.py-- already atnum_stages = 1; only astale
TODOcomment mentions0.fbcode/gem/next_gen-- an ACL'd path, so it lives in the child diff.Phabricator rejects diffs that touch both ACL'd and non-ACL'd paths
(https://fburl.com/no-mixed-paths).
third-party/triton/stableandfbcode/triton_mtia/third_party/triton-- reverted per request; these vendored TLX tutorial copies should be
updated via the upstream sync instead.
pyper_models/*/archive/andminimal_viable_ai/p4p/*/cloned_files/, and unrelatednum_stagesconcepts(dataswarm, shardmanager, eval pipelines).
Note: D114743914 (
[Triton] [Addmm] Fix num_stages range) may overlap; expectto prune anything already covered by other diffs.
Reviewed By: jianyuh
Differential Revision: D114751669