Skip to content

[Triton] Integrate reorder-loop-loads into LoopUnroll pass#775

Open
isaname wants to merge 8 commits into
flagos-ai:mainfrom
isaname:oyrh_dev
Open

[Triton] Integrate reorder-loop-loads into LoopUnroll pass#775
isaname wants to merge 8 commits into
flagos-ai:mainfrom
isaname:oyrh_dev

Conversation

@isaname

@isaname isaname commented Jul 9, 2026 •
edited
Loading

Copy link
Copy Markdown

Background

In Triton kernels, load operations and compute operations are interleaved after loop unrolling. On some hardware, this interleaving prevents memory access latency from being effectively hidden — because each load's result is immediately consumed by the computation that follows it, leaving the pipeline underutilized.

The ideal approach is to hoist (cluster) the load operations from each unrolled iteration to the front, so that multiple loads can be in-flight simultaneously, thereby exploiting memory system bandwidth and hiding latency.

Design

1. Reorder logic integrated into LoopUnroll

For loops with tt.reorder attribute:

  • Full unroll: loads and their transitive deps are clustered to the front of the unrolled op range in the parent block.
  • Partial unroll: same clustering is applied within the expanded loop body.

Key helpers: willFullyUnroll() (static trip count check), collectDepsInBlock() (transitive dep collection), reorderLoadsInRange() (load-first reordering preserving relative order within each group).

2. User interface

Added reorder=False parameter to tl.range. When reorder=True with loop_unroll_factor set, the tt.reorder attribute is attached to the scf.for, triggering load clustering during unroll.

for i in tl.range(0, N, loop_unroll_factor=4, reorder=True):
    a = tl.load(a_ptr + i * BLOCK)
    b = tl.load(b_ptr + i * BLOCK)
    c = a + b
    tl.store(c_ptr + i * BLOCK, c)

Performance

fused_inv_rope_fp8_quant

num_tokens num_heads n_groups heads_per_group tma_aligned_scales Δ%(pipeline) Δ%(reorder)
256 128 8 16 False +9.41% +12.89%
256 128 8 16 True +11.37% +13.76%
512 128 8 16 False +8.04% +12.49%
512 128 8 16 True +12.26% +12.76%
1024 128 8 16 False +11.11% +17.39%
1024 128 8 16 True +14.46% +16.27%
2048 128 8 16 False +13.58% +18.02%
2048 128 8 16 True +16.42% +17.92%
4096 128 8 16 False +12.46% +17.03%
4096 128 8 16 True +15.78% +17.88%

@CLAassistant

CLAassistant commented Jul 9, 2026 •
edited
Loading

Copy link
Copy Markdown

CLA assistant check
All committers have signed the CLA.

@github-actions github-actions Bot added the main label Jul 13, 2026
@zhzhcookie

zhzhcookie commented Jul 13, 2026

Copy link
Copy Markdown
Collaborator

Please refer to PR 763 to add the C++ macro switch.
For every change to the Python code, add the comment # flagtree reorder-loop-loads.

@sunnycase

sunnycase commented Jul 15, 2026

Copy link
Copy Markdown
Collaborator

Thanks for working on this optimization. Could you please avoid modifying the upstream triton.language.range / tl.range primitive? Since reorder is a FlagTree/TLE-specific extension, please introduce tle.range(..., reorder=False) with reorder as a boolean parameter, and keep the existing Triton primitive unchanged.

Would it also be possible to move as much of the backend-specific implementation as possible out of upstream Triton-owned directories and into the TLE subtree? If a hook in an upstream-owned file is unavoidable, please keep it minimal and place the TLE-specific code behind the appropriate TLE build guard. This will minimize divergence from upstream and make ownership of the extension explicit.

Please update the relevant call sites and tests to use tle.range. Thank you.

zhzhcookie
zhzhcookie reviewed Jul 16, 2026
View reviewed changes
Comment thread python/tutorials/12-fused_inv_rope_fp8_quant_reorder.py

@zhzhcookie zhzhcookie Jul 16, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Please add this test to nv3.6-build-and-test.yml

zhzhcookie
zhzhcookie reviewed Jul 16, 2026
View reviewed changes
Comment thread python/triton/compiler/code_generator.py
@isaname
[TLE] Refactor reorder-loop-loads as TLE-specific extension
6ef694f 
Addresses PR feedback by removing `reorder` from upstream `tl.range` and moving
it entirely into the TLE namespace, minimizing upstream Triton modifications.

Frontend changes:
- Remove `reorder` parameter from `triton.language.range` (python/triton/language/core.py)
- Introduce `tle.range(..., reorder=False)` that inherits from `tl.range` and adds
  the `reorder` boolean parameter (python/triton/experimental/tle/language/gpu/core.py)
- Export `tle.range` through TLE language API (__init__.py files)
- Update compiler to recognize `tle.range` iterator (python/triton/compiler/code_generator.py)

Backend changes:
- Extract all reorder-loop-loads implementation (~108 lines) from LoopUnroll.cpp
  into dedicated TLE-owned files:
  - third_party/tle/dialect/include/Transforms/ReorderLoopLoads.h
  - third_party/tle/dialect/lib/Transforms/ReorderLoopLoads.cpp
- Keep minimal hook in LoopUnroll.cpp behind __FLAGTREE_REORDER_LOOP_LOADS__ guard
- Conditionally compile TLE implementation only when FLAGTREE_TLE is enabled
  (lib/Dialect/Triton/Transforms/CMakeLists.txt)

Call sites:
- Update tutorial to use `tle.range(..., reorder=True)` (python/tutorials/12-fused_inv_rope_fp8_quant_reorder.py)

Impact:
- Upstream `tl.range` remains unchanged
- Clear ownership: `reorder` is explicitly a FlagTree/TLE extension
- Minimal upstream footprint: reduces divergence from mainline Triton
@isaname
isaname requested a review from i3wanna2 as a code owner July 21, 2026 11:48
@github-actions github-actions Bot added the tle label Jul 21, 2026
@github-actions github-actions Bot added DOC Improvements or additions to documentation nvidia CORE CI/CD labels Jul 22, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@zhzhcookie zhzhcookie zhzhcookie left review comments
@menchunlei menchunlei Awaiting requested review from menchunlei menchunlei is a code owner
@sunnycase sunnycase Awaiting requested review from sunnycase sunnycase is a code owner
@Galaxy1458 Galaxy1458 Awaiting requested review from Galaxy1458 Galaxy1458 is a code owner
@i3wanna2 i3wanna2 Awaiting requested review from i3wanna2 i3wanna2 is a code owner

At least 1 approving review is required to merge this pull request.

Assignees

No one assigned

Labels

CI/CD CORE DOC Improvements or additions to documentation main nvidia tle triton_v3.6.x

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@isaname @CLAassistant @zhzhcookie @sunnycase