[triton 3.2] [Inductor] `Assertion idx < size() failed.`

[jataylo]

Problem Description

More info https://github.com/pytorch/pytorch/issues/139621

Encountered while testing for preview release/2.6 builds as part of https://github.com/pytorch/pytorch/issues/139175 python test_flex_attention.py -k "test_load_from_bias_seq_only_float16

Traceback

test_load_from_bias_seq_only_float16 (main.TestFlexAttention) ... python: /root/.triton/llvm/llvm-b5cc222d-ubuntu-x64/include/llvm/ADT/SmallVector.h:291: T& llvm::SmallVectorTemplateCommon<T, >::operator[](llvm::SmallVectorTemplateCommon<T, >::size_type) [with T = unsigned int; = void; llvm::SmallVectorTemplateCommon<T, >::reference = unsigned int&; llvm::SmallVectorTemplateCommon<T, >::size_type = long unsigned int]: Assertion `idx < size()' failed.

Failed version reproducer:

• Image used: rocm/pytorch-private:jataylo_pt2.6_preview
• Triton commit: 3889f3f3b97b817741e308c173409927b7c4536f
• Reproducer: https://gist.github.com/jataylo/dfd4179d26e8fd5019a81189949e6f3b

Note this UT passed before the triton upgrade, to reproduce passing version:

• Docker image used: rocm/pytorch-nightly:latest
• Triton branch: release/3.1.x
• Reproducer: https://gist.github.com/jataylo/9ff17e454a84b4297625c917dc8c989d

Operating System

Ubuntu

CPU

.

GPU

MI250X

ROCm Version

ROCm 6.2.0

ROCm Component

No response

Steps to Reproduce

No response

(Optional for Linux users) Output of /opt/rocm/bin/rocminfo --support

No response

Additional Information

No response

[jataylo]

cc: @antiagainst

[jataylo]

Upstream PR is filed that should resolve this https://github.com/triton-lang/triton/pull/5084