[PyTorch Upstream]NaN for FlexAttention with triton xpu

[liangan1]

Describe the bug

There is NaN when running the flexattntion using the triton kernel with XPU but there is no this issue for CUDA backend. You can reproduce this issue by referring to https://jira.devtools.intel.com/browse/TRITONXPU-174

Environment details

Triton commit to reproduce:

commit 650428b (HEAD -> main, origin/main, origin/HEAD)

Author: Pavel Chekin pavel.chekin@intel.com

Date: Mon Dec 16 22:39:04 2024 -0800

Ignore cleanup errors in cache teardown. (#3020)

To avoid PermissionError on Windows where certain .pyd files are locked

[riverliuintel]

PyTorch 2.7 required bug fix

[alexbaden]

I bisected and it appears this is caused by

6588f0de8fb46be7839ed0f111c99b97646f210c is the first bad commit
commit 6588f0de8fb46be7839ed0f111c99b97646f210c
Author: Lukas Sommer <lukas.sommer@codeplay.com>
Date:   Tue Dec 3 13:28:51 2024 +0100

    [XPU][TritonGEN] Replace subgroup reduce and scan (#2893)
    
    Replace usage of the TritonGEN dialect subgroup reduce and scan
    operations with equivalent operations from the SPIR-V dialect.
    
    This closes #2892.
    
    ---------
    
    Signed-off-by: Lukas Sommer <lukas.sommer@codeplay.com>

It looks like the instruction used is _Z27__spirv_GroupNonUniformFMaxiif. I will try to attach the Triton IRs to this ticket. @sommerlukas any ideas?

[alexbaden]

broken_flex_attention.zip

[sommerlukas]

It looks like the instruction used is _Z27__spirv_GroupNonUniformFMaxiif. I will try to attach the Triton IRs to this ticket.

@alexbaden Do you know which function was called instead before my changes?

[alexbaden]

Looks like we go from

  %4732 = call spir_func float @_Z32sub_group_non_uniform_reduce_maxf(float %4676)

to

  %4725 = call spir_func float @_Z27__spirv_GroupNonUniformFMaxiif(i32 3, i32 0, float %4669) #9, !dbg !92

attaching the complete working IR below.

working_flex_attention.zip

[sommerlukas]

That replacement looks correct to me, at least from the SPIR-V spec of GroupNonUniformFMax.

The first parameter is the Scope, in our case 3, which means Subgroup and matches the previous function.

The second parameter is the Group_Operation, in our case 0, which means Reduce (as opposed to ClusteredReduce) and also matches the previous function.

[victor-eds]

👍 to Lukas. I have checked what the working_flex_attention.zip .spv file compiles to with:

ocloc compile -device pvc -o foo  -file ~/triton_tem_fused_0.spv -spirv_input

and it's indeed the same SPIR-V builtin. There are some IR differences as you note, tho.

[alexbaden]

will be fixed by llvm/llvm-project#122131

The issue appears to be the lack of the convergent function attribute on the spirv group reduce instructions coming from the MLIR SPIRV dialect. Although the non-uniform spirv instructions don't require convergent control flow, the lack of the convergent attribute allows LLVM to wildly alter the control flow (e.g. by duplicating and sinking the group reduced instructions) which breaks the program at later stages in the lowering. Once the LLVM PR is merged we will need OpenAI to bump the pin; or as a last resort we can create our own LLVM build again.

[liangan1]

Can you specify a Triton commit to validate the functionality of this fix?

[sommerlukas]

Can you specify a Triton commit to validate the functionality of this fix?

Hi @liangan1, we first need the upstream LLVM PR that Alex linked to be merged. Then upstream Triton needs to update their LLVM base commit (they usually do that on a (bi-)weekly basis) and then we get the updated LLVM base commit through a pulldown of upstream Triton into our project.

We can ping here once that's done.