Skip to content

Replace num_stages=0 with num_stages=1 - #231

Open
njriasan wants to merge 1 commit into
mainfrom
export-D114751669
Open

njriasan wants to merge 1 commit into
mainfrom
export-D114751669

Conversation

@njriasan

@njriasan njriasan commented Aug 4, 2026

Copy link
Copy Markdown

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

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
@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label Aug 4, 2026
@meta-codesync

meta-codesync Bot commented Aug 4, 2026

Copy link
Copy Markdown
Contributor

@njriasan has exported this pull request. If you are a Meta employee, you can view the originating Diff in D114751669.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CLA Signed This label is managed by the Meta Open Source bot. meta-exported

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant