[FlexAttenion] Crash When Number of Variables Exceeds Assigned Values in Multi-Assignment

[retonym]

Describe the bug

Issue Description:

While testing the PyTorch flex-attention module, we encountered a scenario where the number of variables does not match the number of assigned values in a multi-assignment statement. For example:

a, b, c = 1, 1

Although this code is syntactically unusual , it executes successfully with the public Triton on Nvidia GPU. However, when running on Intel Triton, it fails with the error message:

IndexError('list index out of range')

We would like to request support for handling such cases in Intel Triton. Thank you for your assistance.

Steps to Reproduce:

Below is a minimal reproducible example that demonstrates the issue:

import triton
import triton.language as tl

import torch

device = 'xpu' if torch.xpu.is_available() else 'cuda'

@triton.jit
def test_varlen_list(A, BLOCK_SIZE: tl.constexpr):
  pid = tl.program_id(0)

  block_start = pid * BLOCK_SIZE

  offsets = block_start + tl.arange(0, BLOCK_SIZE)

  a, b, c = 1, 1

  tl.store(A + offsets, a)


BLOCK_SIZE = 32
TENSOR_SIZE = 512

A = torch.empty(TENSOR_SIZE, dtype=torch.float32, device=device)

grid = lambda meta: ((A.numel() + BLOCK_SIZE - 1) // BLOCK_SIZE,)

test_varlen_list[grid](A, BLOCK_SIZE)
print(A)

Environment details

PyTorch commit: 520079b986b76dc42e7fec3f992da3b1771e3192
Intel Triton commit: b7840ba
Public Triton commit:4b3bb1f8

[liangan1]

cc @EikanWang @riverliuintel

[quintinwang5]

According to my tests, this code is just a mistake. It works by accident(previous frontend did not have strict checks). On CUDA, with public triton commit: 4b3bb1f8, the reproducer works, but c is undefined(trying to store c into A triggers an undefined error of c). With the newest commit 74fe0c1 until now (Feb. 16), the result is:

Traceback (most recent call last):
  File "/home/pt-gpu/tmp/x.py", line 28, in <module>
    test_varlen_list[grid](A, BLOCK_SIZE)
    ~~~~~~~~~~~~~~~~~~~~~~^^^^^^^^^^^^^^^
  File "/home/pt-gpu/miniforge3/envs/qt_tt/lib/python3.13/site-packages/triton/runtime/jit.py", line 336, in <lambda>
    return lambda *args, **kwargs: self.run(grid=grid, warmup=False, *args, **kwargs)
                                   ~~~~~~~~^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/pt-gpu/miniforge3/envs/qt_tt/lib/python3.13/site-packages/triton/runtime/jit.py", line 563, in run
    kernel = self.compile(src, target=target, options=options.__dict__)
  File "/home/pt-gpu/miniforge3/envs/qt_tt/lib/python3.13/site-packages/triton/compiler/compiler.py", line 277, in compile
    module = src.make_ir(options, codegen_fns, module_map, context)
  File "/home/pt-gpu/miniforge3/envs/qt_tt/lib/python3.13/site-packages/triton/compiler/compiler.py", line 81, in make_ir
    return ast_to_ttir(self.fn, self, context=context, options=options, codegen_fns=codegen_fns,
                       module_map=module_map)
triton.compiler.errors.CompilationError: at 8:2:
def test_varlen_list(A, BLOCK_SIZE: tl.constexpr):
  pid = tl.program_id(0)

  block_start = pid * BLOCK_SIZE

  offsets = block_start + tl.arange(0, BLOCK_SIZE)

  a, b, c = 1, 1
  ^
IndexError('list index out of range')

This error occurs during frontend stage ast_to_ttir. Intel Triton fails because we merged more frontend changes in compiler.py/ code_generator.py.

[retonym]

Thanks to @quintinwang5 for the investigation. Let's observe how Meta addresses this issue after their upgrade of the Triton commit.