Skip to content

[Fix] Avoid rebinding chunk loop variables for int64 anchoring - #1188

Merged
zhiyuan1i merged 1 commit into
mainfrom
fix/int64-loop-rebind
Aug 28, 2026
Merged

zhiyuan1i merged 1 commit into
mainfrom
fix/int64-loop-rebind

Conversation

@zhiyuan1i

Copy link
Copy Markdown
Collaborator

Summary

Triton shipped with PyTorch 2.7 rejects re-assigning a loop induction variable with a different type:

AssertionError: Loop-carried variable i_t has initial type int32 but is re-assigned to int64 in loop!

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 fresh i_t_int64 variable instead, matching the convention already used in chunk_delta_h.py. Semantics unchanged.

Test plan

  • Unit tests added/modified: none (no behavior change).
  • Dependent tests run: python scripts/find_dependent_tests.py over the 5 changed files (30 test files across ops/layers/context_parallel).
    • torch 2.7.1 + triton 3.3.1 (same versions as the failing CI env): pre-fix reproduces the exact CI failure (test_precond_gated_delta.py::test_chunk, same failing case, 8 passed / 1 failed); post-fix the full file passes (38/38).
    • torch 2.10 (newer Triton): test_ttt.py, test_linear_attn.py, test_retention.py pass.
    • The full 30-file dependent suite on torch 2.7.1 is running; this PR's own PyTorch 2.7 CI job covers the same scope.
  • Varlen / CP / model tests: the dependent list includes tests/context_parallel and varlen cases; all exercised by CI.

Benchmark / NCU (kernel changes only)

Neutral — pure variable rename, no performance-related change.

Breaking changes

None.

Checklist

  • I have read CONTRIBUTING.md and follow its conventions (code style, docstrings, commit prefixes).
  • I have read AGENTS.md and, where my change matches its scope, the relevant skill under .agents/skills.
  • Dependent tests pass locally or in CI, and new behavior is covered by tests where applicable (tick as N/A for changes with no testable code, e.g. docs-only).
  • Kernel changes include same-hardware before/after benchmark numbers, dense + varlen where applicable (tick as N/A when no kernel code changed).
  • This PR is minor/cosmetic-only (typo, formatting, style-only tweaks) — tick only if it is, and justify below.

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.
@zhiyuan1i zhiyuan1i added the bug Something isn't working label Aug 28, 2026
@zhiyuan1i
zhiyuan1i merged commit bccaf2d into main Aug 28, 2026
11 of 13 checks passed
@zhiyuan1i
zhiyuan1i deleted the fix/int64-loop-rebind branch August 28, 2026 08:39
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bug Something isn't working

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant