[Fix] Avoid rebinding chunk loop variables for int64 anchoring - #1188
Merged
Merged
Conversation
Triton shipped with PyTorch 2.7 rejects re-assigning a loop induction variable with a different type (initial type int32 but is re-assigned to int64 in loop), breaking the H100 PyTorch 2.7 ops job since #1173. Use a fresh i_t_int64 variable instead, matching the existing convention in chunk_delta_h.py. Semantics unchanged.
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
Triton shipped with PyTorch 2.7 rejects re-assigning a loop induction variable with a different type:
Since #1173, eleven chunk kernels cast the loop variable in place (
i_t = i_t.to(tl.int64)), so ops using them fail to compile under the H100 PyTorch 2.7 CI env (seen on #1184 and #1111, e.g.tests/ops/test_precond_gated_delta.py::test_chunk). Use a freshi_t_int64variable instead, matching the convention already used inchunk_delta_h.py. Semantics unchanged.Test plan
python scripts/find_dependent_tests.pyover the 5 changed files (30 test files across ops/layers/context_parallel).test_precond_gated_delta.py::test_chunk, same failing case, 8 passed / 1 failed); post-fix the full file passes (38/38).test_ttt.py,test_linear_attn.py,test_retention.pypass.tests/context_paralleland varlen cases; all exercised by CI.Benchmark / NCU (kernel changes only)
Neutral — pure variable rename, no performance-related change.
Breaking changes
None.
Checklist