[Bug] triton error when run chunk gated deltanet op

[LouChao98]

Checklist

• I have checked FAQs and existing issues for similar problems
• Please report this bug in English to ensure wider understanding and support

Describe the Bug

Env: triton==3.2.0, H800

When run into

flash-linear-attention/fla/ops/common/chunk_delta_h.py

Line 27 in fec1091

def chunk_gated_delta_rule_fwd_kernel_h(

I get

python3: /project/lib/Tools/LinearLayout.cpp:562: mlir::triton::LinearLayout mlir::triton::LinearLayout::reshapeOuts(llvm::ArrayRef<std::pair<mlir::StringAttr, int> >) const: Assertion `getTotalOutDimSize() == std::accumulate( newOutDims.begin(), newOutDims.end(), 1, [&](int32_t acc, auto &outDim) { return acc * outDim.second; })' failed.

Using triton==3.1.0 does not encouter this error.

Steps to Reproduce the Bug

import torch
from fla.ops.common.chunk_delta_h import chunk_gated_delta_rule_fwd_h

k = torch.randn(1, 8192, 36, 256, device="cuda", dtype=torch.bfloat16)
w = torch.randn(1, 8192, 36, 256, device="cuda", dtype=torch.bfloat16)
u = torch.randn(1, 8192, 36, 256, device="cuda", dtype=torch.bfloat16)
g = torch.randn(1, 8192, 36, device="cuda", dtype=torch.float32)
head_first = False
chunk_size = 64

chunk_gated_delta_rule_fwd_h(k=k, w=w, u=u, g=g, head_first=head_first, chunk_size=chunk_size)

Expected Behavior

no error

Environment Information

1. Torch: 2.3.1
2. Triton: 3.2.0

[sustcsonglin]

This is an known issue triton-lang/triton#5609
Either reinstall triton from scratch or use Triton3.1

[LouChao98]

I installed Triton from scratch, and now the version is 3.2.0+gitc1ed673b. The same error is raised, which I've posted below.
After removing num_wraps=8 and num_wraps=16 from the autotune configuration, the error disappeared.

So, this is a triton bug and I should report to them, right? The test case in the linked triton issue passes on my env. Something is different. But I don't know how to make a minimum reproduce code...

python: /workspace/triton/lib/Dialect/TritonGPU/IR/LinearLayoutConversions.cpp:1131: mlir::triton::LinearLayout mlir::triton::gpu::{anonymous}::chooseStMatrixLayoutLeadingOffset(mlir::MLIRContext*, mlir::RankedTensorType, int): Assertion `instrN >= numColsPerChunk && "Each chunk is filled in with a single warp"' failed.
#blocked = #ttg.blocked<{sizePerThread = [1, 8], threadsPerWarp = [1, 32], warpsPerCTA = [8, 1], order = [1, 0]}>
#blocked1 = #ttg.blocked<{sizePerThread = [1, 8], threadsPerWarp = [4, 8], warpsPerCTA = [8, 1], order = [1, 0]}>
#blocked2 = #ttg.blocked<{sizePerThread = [8, 1], threadsPerWarp = [32, 1], warpsPerCTA = [1, 8], order = [0, 1]}>
#blocked3 = #ttg.blocked<{sizePerThread = [1], threadsPerWarp = [32], warpsPerCTA = [8], order = [0]}>
#mma = #ttg.nvidia_mma<{versionMajor = 3, versionMinor = 0, warpsPerCTA = [8, 1], instrShape = [16, 64, 16]}>
#mma1 = #ttg.nvidia_mma<{versionMajor = 3, versionMinor = 0, warpsPerCTA = [8, 1], instrShape = [16, 32, 16]}>
#shared = #ttg.nvmma_shared<{swizzlingByteWidth = 128, transposed = false, elementBitWidth = 16}>
#shared1 = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0]}>
#shared2 = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0, 1]}>
#shared3 = #ttg.swizzled_shared<{vec = 8, perPhase = 1, maxPhase = 8, order = [1, 0]}>
#shared4 = #ttg.swizzled_shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [1, 0]}>
#shared5 = #ttg.swizzled_shared<{vec = 8, perPhase = 1, maxPhase = 8, order = [0, 1]}>
#smem = #ttg.shared_memory
module attributes {"ttg.num-ctas" = 1 : i32, "ttg.num-warps" = 8 : i32, ttg.target = "cuda:90", "ttg.threads-per-warp" = 32 : i32} {
tt.func public @chunk_gated_delta_rule_fwd_kernel_h(%arg0: !tt.ptr {tt.divisibility = 16 : i32}, %arg1: !tt.ptr {tt.divisibility = 16 : i32}, %arg2: !tt.ptr {tt.divisibility = 16 : i32}, %arg3: !tt.ptr {tt.divisibility = 16 : i32}, %arg4: !tt.ptr {tt.divisibility = 16 : i32}, %arg5: !tt.ptr {tt.divisibility = 16 : i32}, %arg6: i32) attributes {noinline = false} {
%cst = arith.constant dense<0.000000e+00> : tensor<256x64xf32, #mma>
%c0_i32 = arith.constant 0 : i32
%c65536_i32 = arith.constant 65536 : i32
%c63_i32 = arith.constant 63 : i32
%c36_i32 = arith.constant 36 : i32
%c1_i32 = arith.constant 1 : i32
%c256_i32 = arith.constant 256 : i32
%c64_i32 = arith.constant 64 : i32
%cst_0 = arith.constant dense<256> : tensor<1x256xi64, #blocked>
%cst_1 = arith.constant dense<0> : tensor<1x256xi64, #blocked>
%cst_2 = arith.constant dense<0> : tensor<64x1xi64, #blocked>
%cst_3 = arith.constant dense<0> : tensor<64x1xi64, #blocked1>
%cst_4 = arith.constant dense<9216> : tensor<64x1xi64, #blocked>
%cst_5 = arith.constant dense<9216> : tensor<64x1xi64, #blocked1>
%cst_6 = arith.constant dense<9216> : tensor<1x64xi64, #blocked2>
%c2_i32 = arith.constant 2 : i32
%false = arith.constant false
%cst_7 = arith.constant dense<0.000000e+00> : tensor<64x64xf32, #mma1>
%cst_8 = arith.constant dense<256> : tensor<256x1xi64, #blocked2>
%cst_9 = arith.constant dense<256> : tensor<256x1xi64, #blocked1>
%cst_10 = arith.constant dense<0> : tensor<256x1xi64, #blocked2>
%cst_11 = arith.constant dense<0> : tensor<256x1xi64, #blocked1>
%cst_12 = arith.constant dense<0> : tensor<1x64xi64, #blocked2>
%cst_13 = arith.constant dense<0> : tensor<1x64xi64, #blocked1>
%cst_14 = arith.constant dense<256> : tensor<1x64xi64, #blocked1>
%cst_15 = arith.constant dense<36> : tensor<64xi64, #blocked3>
%cst_16 = arith.constant dense<0> : tensor<64xi64, #blocked3>
%0 = tt.get_program_id x : i32
%1 = tt.get_program_id y : i32
%2 = tt.get_program_id z : i32
%3 = arith.divsi %2, %c36_i32 : i32
%4 = arith.remsi %2, %c36_i32 : i32
%5 = arith.muli %3, %arg6 : i32
%6 = arith.addi %arg6, %c63_i32 : i32
%7 = arith.divsi %6, %c64_i32 : i32
%8 = arith.muli %3, %7 : i32
%9 = arith.muli %0, %c256_i32 : i32
%10 = arith.muli %1, %c64_i32 : i32
%11 = arith.extsi %9 : i32 to i64
%12 = arith.extsi %10 : i32 to i64
%13 = tt.splat %11 : i64 -> tensor<256xi64, #ttg.slice<{dim = 1, parent = #blocked1}>>
%14 = tt.splat %11 : i64 -> tensor<256xi64, #ttg.slice<{dim = 1, parent = #blocked2}>>
%15 = tt.splat %11 : i64 -> tensor<256xi64, #ttg.slice<{dim = 0, parent = #blocked}>>
%16 = tt.make_range {end = 256 : i32, start = 0 : i32} : tensor<256xi32, #ttg.slice<{dim = 1, parent = #blocked1}>>
%17 = tt.make_range {end = 256 : i32, start = 0 : i32} : tensor<256xi32, #ttg.slice<{dim = 1, parent = #blocked2}>>
%18 = tt.make_range {end = 256 : i32, start = 0 : i32} : tensor<256xi32, #ttg.slice<{dim = 0, parent = #blocked}>>
%19 = arith.extsi %16 : tensor<256xi32, #ttg.slice<{dim = 1, parent = #blocked1}>> to tensor<256xi64, #ttg.slice<{dim = 1, parent = #blocked1}>>
%20 = arith.extsi %17 : tensor<256xi32, #ttg.slice<{dim = 1, parent = #blocked2}>> to tensor<256xi64, #ttg.slice<{dim = 1, parent = #blocked2}>>
%21 = arith.extsi %18 : tensor<256xi32, #ttg.slice<{dim = 0, parent = #blocked}>> to tensor<256xi64, #ttg.slice<{dim = 0, parent = #blocked}>>
%22 = arith.addi %13, %19 : tensor<256xi64, #ttg.slice<{dim = 1, parent = #blocked1}>>
%23 = arith.addi %14, %20 : tensor<256xi64, #ttg.slice<{dim = 1, parent = #blocked2}>>
%24 = arith.addi %15, %21 : tensor<256xi64, #ttg.slice<{dim = 0, parent = #blocked}>>
%25 = tt.expand_dims %22 {axis = 1 : i32} : tensor<256xi64, #ttg.slice<{dim = 1, parent = #blocked1}>> -> tensor<256x1xi64, #blocked1>
%26 = tt.expand_dims %23 {axis = 1 : i32} : tensor<256xi64, #ttg.slice<{dim = 1, parent = #blocked2}>> -> tensor<256x1xi64, #blocked2>
%27 = arith.muli %25, %cst_9 : tensor<256x1xi64, #blocked1>
%28 = tt.broadcast %27 : tensor<256x1xi64, #blocked1> -> tensor<256x64xi64, #blocked1>
%29 = tt.splat %12 : i64 -> tensor<64xi64, #ttg.slice<{dim = 0, parent = #blocked1}>>
%30 = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32, #ttg.slice<{dim = 0, parent = #blocked1}>>
%31 = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32, #blocked3>
%32 = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32, #ttg.slice<{dim = 0, parent = #blocked2}>>
%33 = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32, #ttg.slice<{dim = 1, parent = #blocked}>>
%34 = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32, #ttg.slice<{dim = 1, parent = #blocked1}>>
%35 = arith.extsi %30 : tensor<64xi32, #ttg.slice<{dim = 0, parent = #blocked1}>> to tensor<64xi64, #ttg.slice<{dim = 0, parent = #blocked1}>>
%36 = arith.extsi %31 : tensor<64xi32, #blocked3> to tensor<64xi64, #blocked3>
%37 = arith.extsi %32 : tensor<64xi32, #ttg.slice<{dim = 0, parent = #blocked2}>> to tensor<64xi64, #ttg.slice<{dim = 0, parent = #blocked2}>>
%38 = arith.extsi %33 : tensor<64xi32, #ttg.slice<{dim = 1, parent = #blocked}>> to tensor<64xi64, #ttg.slice<{dim = 1, parent = #blocked}>>
%39 = arith.extsi %34 : tensor<64xi32, #ttg.slice<{dim = 1, parent = #blocked1}>> to tensor<64xi64, #ttg.slice<{dim = 1, parent = #blocked1}>>
%40 = arith.addi %29, %35 : tensor<64xi64, #ttg.slice<{dim = 0, parent = #blocked1}>>
%41 = tt.expand_dims %40 {axis = 0 : i32} : tensor<64xi64, #ttg.slice<{dim = 0, parent = #blocked1}>> -> tensor<1x64xi64, #blocked1>
%42 = tt.broadcast %41 : tensor<1x64xi64, #blocked1> -> tensor<256x64xi64, #blocked1>
%43 = arith.addi %28, %42 : tensor<256x64xi64, #blocked1>
%44 = arith.cmpi sge, %25, %cst_11 : tensor<256x1xi64, #blocked1>
%45 = arith.cmpi sge, %26, %cst_10 : tensor<256x1xi64, #blocked2>
%46 = arith.cmpi slt, %25, %cst_9 : tensor<256x1xi64, #blocked1>
%47 = arith.cmpi slt, %26, %cst_8 : tensor<256x1xi64, #blocked2>
%48 = arith.andi %44, %46 : tensor<256x1xi1, #blocked1>
%49 = arith.andi %45, %47 : tensor<256x1xi1, #blocked2>
%50 = tt.broadcast %48 : tensor<256x1xi1, #blocked1> -> tensor<256x64xi1, #blocked1>
%51 = tt.broadcast %49 : tensor<256x1xi1, #blocked2> -> tensor<256x64xi1, #blocked2>
%52 = arith.cmpi sge, %41, %cst_13 : tensor<1x64xi64, #blocked1>
%53 = arith.cmpi slt, %41, %cst_14 : tensor<1x64xi64, #blocked1>
%54 = arith.andi %52, %53 : tensor<1x64xi1, #blocked1>
%55 = tt.broadcast %54 : tensor<1x64xi1, #blocked1> -> tensor<256x64xi1, #blocked1>
%56 = arith.andi %50, %55 : tensor<256x64xi1, #blocked1>
%57 = arith.muli %5, %c36_i32 : i32
%58 = tt.addptr %arg4, %57 : !tt.ptr, i32
%59 = arith.addi %57, %4 : i32
%60 = arith.muli %59, %c256_i32 : i32
%61 = tt.addptr %arg0, %60 : !tt.ptr, i32
%62 = arith.extsi %arg6 : i32 to i64
%63 = tt.addptr %arg2, %60 : !tt.ptr, i32
%64 = tt.addptr %arg1, %60 : !tt.ptr, i32
%65 = tt.addptr %arg3, %60 : !tt.ptr, i32
%66 = tt.addptr %58, %4 : !tt.ptr, i32
%67 = tt.splat %66 : !tt.ptr -> tensor<64x!tt.ptr, #blocked3>
%68 = tt.splat %62 : i64 -> tensor<64xi64, #blocked3>
%69 = tt.splat %61 : !tt.ptr -> tensor<256x64x!tt.ptr, #blocked2>
%70 = tt.broadcast %26 : tensor<256x1xi64, #blocked2> -> tensor<256x64xi64, #blocked2>
%71 = tt.splat %62 : i64 -> tensor<1x64xi64, #blocked2>
%72 = tt.splat %63 : !tt.ptr -> tensor<64x256x!tt.ptr, #blocked>
%73 = tt.expand_dims %24 {axis = 0 : i32} : tensor<256xi64, #ttg.slice<{dim = 0, parent = #blocked}>> -> tensor<1x256xi64, #blocked>
%74 = tt.broadcast %73 : tensor<1x256xi64, #blocked> -> tensor<64x256xi64, #blocked>
%75 = tt.splat %62 : i64 -> tensor<64x1xi64, #blocked>
%76 = tt.splat %62 : i64 -> tensor<64x1xi64, #blocked1>
%77 = arith.cmpi sge, %73, %cst_1 : tensor<1x256xi64, #blocked>
%78 = arith.cmpi slt, %73, %cst_0 : tensor<1x256xi64, #blocked>
%79 = arith.andi %77, %78 : tensor<1x256xi1, #blocked>
%80 = tt.broadcast %79 : tensor<1x256xi1, #blocked> -> tensor<64x256xi1, #blocked>
%81 = tt.splat %64 : !tt.ptr -> tensor<64x64x!tt.ptr, #blocked1>
%82 = tt.broadcast %41 : tensor<1x64xi64, #blocked1> -> tensor<64x64xi64, #blocked1>
%83 = tt.broadcast %54 : tensor<1x64xi1, #blocked1> -> tensor<64x64xi1, #blocked1>
%84 = tt.splat %65 : !tt.ptr -> tensor<64x64x!tt.ptr, #blocked1>
%85 = scf.for %arg7 = %c0_i32 to %7 step %c1_i32 iter_args(%arg8 = %cst) -> (tensor<256x64xf32, #mma>) : i32 {
%86 = arith.addi %8, %arg7 : i32
%87 = arith.muli %86, %c36_i32 : i32
%88 = arith.addi %87, %4 : i32
%89 = arith.muli %88, %c65536_i32 : i32
%90 = tt.addptr %arg5, %89 : !tt.ptr, i32
%91 = arith.truncf %arg8 : tensor<256x64xf32, #mma> to tensor<256x64xbf16, #mma>
%92 = ttg.local_alloc %91 : (tensor<256x64xbf16, #mma>) -> !ttg.memdesc<256x64xbf16, #shared, #smem>
%93 = tt.splat %90 : !tt.ptr -> tensor<256x64x!tt.ptr, #blocked1>
%94 = tt.addptr %93, %43 : tensor<256x64x!tt.ptr, #blocked1>, tensor<256x64xi64, #blocked1>
%95 = ttg.convert_layout %91 : tensor<256x64xbf16, #mma> -> tensor<256x64xbf16, #blocked1>
tt.store %94, %95, %56 : tensor<256x64x!tt.ptr, #blocked1>
%96 = arith.addi %arg7, %c1_i32 : i32
%97 = arith.muli %96, %c64_i32 : i32
%98 = arith.minsi %97, %arg6 : i32
%99 = arith.subi %98, %c1_i32 : i32
%100 = arith.muli %99, %c36_i32 : i32
%101 = tt.addptr %58, %100 : !tt.ptr, i32
%102 = tt.addptr %101, %4 : !tt.ptr, i32
%103 = tt.load %102 : !tt.ptr
%104 = arith.muli %arg7, %c64_i32 : i32
%105 = arith.subi %arg6, %104 : i32
%106 = arith.minsi %105, %c64_i32 : i32
%107 = arith.addi %106, %c63_i32 : i32
%108 = arith.divsi %107, %c64_i32 : i32
%109 = tt.splat %103 : f32 -> tensor<64xf32, #ttg.slice<{dim = 0, parent = #blocked2}>>
%110 = ttg.local_alloc : () -> !ttg.memdesc<2x64xf32, #shared1, #smem, mutable>
%111 = ttg.local_alloc : () -> !ttg.memdesc<2x64xf32, #shared1, #smem, mutable>
%112 = ttg.local_alloc : () -> !ttg.memdesc<2x256x64xbf16, #shared2, #smem, mutable>
%113 = ttg.local_alloc : () -> !ttg.memdesc<2x64x256xbf16, #shared3, #smem, mutable>
%114 = ttg.local_alloc : () -> !ttg.memdesc<2x64x64xbf16, #shared4, #smem, mutable>
%115 = arith.cmpi sgt, %108, %c0_i32 : i32
%116 = arith.extsi %104 : i32 to i64
%117 = tt.splat %116 : i64 -> tensor<64xi64, #blocked3>
%118 = tt.splat %116 : i64 -> tensor<64xi64, #ttg.slice<{dim = 0, parent = #blocked2}>>
%119 = tt.splat %116 : i64 -> tensor<64xi64, #ttg.slice<{dim = 1, parent = #blocked}>>
%120 = tt.splat %116 : i64 -> tensor<64xi64, #ttg.slice<{dim = 1, parent = #blocked1}>>
%121 = arith.addi %117, %36 : tensor<64xi64, #blocked3>
%122 = arith.addi %118, %37 : tensor<64xi64, #ttg.slice<{dim = 0, parent = #blocked2}>>
%123 = arith.addi %119, %38 : tensor<64xi64, #ttg.slice<{dim = 1, parent = #blocked}>>
%124 = arith.addi %120, %39 : tensor<64xi64, #ttg.slice<{dim = 1, parent = #blocked1}>>
%125 = arith.muli %121, %cst_15 : tensor<64xi64, #blocked3>
%126 = tt.addptr %67, %125 : tensor<64x!tt.ptr, #blocked3>, tensor<64xi64, #blocked3>
%127 = arith.cmpi sge, %121, %cst_16 : tensor<64xi64, #blocked3>
%128 = arith.cmpi slt, %121, %68 : tensor<64xi64, #blocked3>
%129 = arith.andi %127, %128 : tensor<64xi1, #blocked3>
%130 = ttg.memdesc_subview %110[%c0_i32, %c0_i32] : !ttg.memdesc<2x64xf32, #shared1, #smem, mutable> -> !ttg.memdesc<64xf32, #shared1, #smem, mutable, 2x64>
%131 = tt.splat %115 : i1 -> tensor<64xi1, #blocked3>
%132 = arith.andi %131, %129 : tensor<64xi1, #blocked3>
%133 = ttg.async_copy_global_to_local %126, %130 mask %132 : tensor<64x!tt.ptr, #blocked3> -> <64xf32, #shared1, #smem, mutable, 2x64>
%134 = ttg.async_commit_group %133
%135 = ttg.memdesc_subview %111[%c0_i32, %c0_i32] : !ttg.memdesc<2x64xf32, #shared1, #smem, mutable> -> !ttg.memdesc<64xf32, #shared1, #smem, mutable, 2x64>
%136 = ttg.async_copy_global_to_local %126, %135 mask %132 : tensor<64x!tt.ptr, #blocked3> -> <64xf32, #shared1, #smem, mutable, 2x64>
%137 = ttg.async_commit_group %136
%138 = tt.expand_dims %122 {axis = 0 : i32} : tensor<64xi64, #ttg.slice<{dim = 0, parent = #blocked2}>> -> tensor<1x64xi64, #blocked2>
%139 = arith.muli %138, %cst_6 : tensor<1x64xi64, #blocked2>
%140 = tt.broadcast %139 : tensor<1x64xi64, #blocked2> -> tensor<256x64xi64, #blocked2>
%141 = arith.addi %70, %140 : tensor<256x64xi64, #blocked2>
%142 = tt.addptr %69, %141 : tensor<256x64x!tt.ptr, #blocked2>, tensor<256x64xi64, #blocked2>
%143 = arith.cmpi sge, %138, %cst_12 : tensor<1x64xi64, #blocked2>
%144 = arith.cmpi slt, %138, %71 : tensor<1x64xi64, #blocked2>
%145 = arith.andi %143, %144 : tensor<1x64xi1, #blocked2>
%146 = tt.broadcast %145 : tensor<1x64xi1, #blocked2> -> tensor<256x64xi1, #blocked2>
%147 = arith.andi %51, %146 : tensor<256x64xi1, #blocked2>
%148 = ttg.memdesc_subview %112[%c0_i32, %c0_i32, %c0_i32] : !ttg.memdesc<2x256x64xbf16, #shared2, #smem, mutable> -> !ttg.memdesc<256x64xbf16, #shared2, #smem, mutable, 2x256x64>
%149 = tt.splat %115 : i1 -> tensor<256x64xi1, #blocked2>
%150 = arith.andi %149, %147 : tensor<256x64xi1, #blocked2>
%151 = ttg.async_copy_global_to_local %142, %148 mask %150 : tensor<256x64x!tt.ptr, #blocked2> -> <256x64xbf16, #shared2, #smem, mutable, 2x256x64>
%152 = ttg.async_commit_group %151
%153 = tt.expand_dims %123 {axis = 1 : i32} : tensor<64xi64, #ttg.slice<{dim = 1, parent = #blocked}>> -> tensor<64x1xi64, #blocked>
%154 = tt.expand_dims %124 {axis = 1 : i32} : tensor<64xi64, #ttg.slice<{dim = 1, parent = #blocked1}>> -> tensor<64x1xi64, #blocked1>
%155 = arith.muli %153, %cst_4 : tensor<64x1xi64, #blocked>
%156 = arith.muli %154, %cst_5 : tensor<64x1xi64, #blocked1>
%157 = tt.broadcast %155 : tensor<64x1xi64, #blocked> -> tensor<64x256xi64, #blocked>
%158 = arith.addi %157, %74 : tensor<64x256xi64, #blocked>
%159 = tt.addptr %72, %158 : tensor<64x256x!tt.ptr, #blocked>, tensor<64x256xi64, #blocked>
%160 = arith.cmpi sge, %153, %cst_2 : tensor<64x1xi64, #blocked>
%161 = arith.cmpi sge, %154, %cst_3 : tensor<64x1xi64, #blocked1>
%162 = arith.cmpi slt, %153, %75 : tensor<64x1xi64, #blocked>
%163 = arith.cmpi slt, %154, %76 : tensor<64x1xi64, #blocked1>
%164 = arith.andi %160, %162 : tensor<64x1xi1, #blocked>
%165 = arith.andi %161, %163 : tensor<64x1xi1, #blocked1>
%166 = tt.broadcast %164 : tensor<64x1xi1, #blocked> -> tensor<64x256xi1, #blocked>
%167 = arith.andi %166, %80 : tensor<64x256xi1, #blocked>
%168 = ttg.memdesc_subview %113[%c0_i32, %c0_i32, %c0_i32] : !ttg.memdesc<2x64x256xbf16, #shared3, #smem, mutable> -> !ttg.memdesc<64x256xbf16, #shared3, #smem, mutable, 2x64x256>
%169 = tt.splat %115 : i1 -> tensor<64x256xi1, #blocked>
%170 = arith.andi %169, %167 : tensor<64x256xi1, #blocked>
%171 = ttg.async_copy_global_to_local %159, %168 mask %170 : tensor<64x256x!tt.ptr, #blocked> -> <64x256xbf16, #shared3, #smem, mutable, 2x64x256>
%172 = ttg.async_commit_group %171
%173 = tt.broadcast %156 : tensor<64x1xi64, #blocked1> -> tensor<64x64xi64, #blocked1>
%174 = arith.addi %173, %82 : tensor<64x64xi64, #blocked1>
%175 = tt.addptr %81, %174 : tensor<64x64x!tt.ptr, #blocked1>, tensor<64x64xi64, #blocked1>
%176 = tt.broadcast %165 : tensor<64x1xi1, #blocked1> -> tensor<64x64xi1, #blocked1>
%177 = arith.andi %176, %83 : tensor<64x64xi1, #blocked1>
%178 = ttg.memdesc_subview %114[%c0_i32, %c0_i32, %c0_i32] : !ttg.memdesc<2x64x64xbf16, #shared4, #smem, mutable> -> !ttg.memdesc<64x64xbf16, #shared4, #smem, mutable, 2x64x64>
%179 = tt.splat %115 : i1 -> tensor<64x64xi1, #blocked1>
%180 = arith.andi %179, %177 : tensor<64x64xi1, #blocked1>
%181 = ttg.async_copy_global_to_local %175, %178 mask %180 : tensor<64x64x!tt.ptr, #blocked1> -> <64x64xbf16, #shared4, #smem, mutable, 2x64x64>
%182 = ttg.async_commit_group %181
%183 = arith.cmpi sgt, %108, %c1_i32 : i32
%184 = arith.addi %104, %c64_i32 : i32
%185 = arith.extsi %184 : i32 to i64
%186 = tt.splat %185 : i64 -> tensor<64xi64, #blocked3>
%187 = tt.splat %185 : i64 -> tensor<64xi64, #ttg.slice<{dim = 0, parent = #blocked2}>>
%188 = tt.splat %185 : i64 -> tensor<64xi64, #ttg.slice<{dim = 1, parent = #blocked}>>
%189 = tt.splat %185 : i64 -> tensor<64xi64, #ttg.slice<{dim = 1, parent = #blocked1}>>
%190 = arith.addi %186, %36 : tensor<64xi64, #blocked3>
%191 = arith.addi %187, %37 : tensor<64xi64, #ttg.slice<{dim = 0, parent = #blocked2}>>
%192 = arith.addi %188, %38 : tensor<64xi64, #ttg.slice<{dim = 1, parent = #blocked}>>
%193 = arith.addi %189, %39 : tensor<64xi64, #ttg.slice<{dim = 1, parent = #blocked1}>>
%194 = arith.muli %190, %cst_15 : tensor<64xi64, #blocked3>
%195 = tt.addptr %67, %194 : tensor<64x!tt.ptr, #blocked3>, tensor<64xi64, #blocked3>
%196 = arith.cmpi sge, %190, %cst_16 : tensor<64xi64, #blocked3>
%197 = arith.cmpi slt, %190, %68 : tensor<64xi64, #blocked3>
%198 = arith.andi %196, %197 : tensor<64xi1, #blocked3>
%199 = ttg.memdesc_subview %110[%c1_i32, %c0_i32] : !ttg.memdesc<2x64xf32, #shared1, #smem, mutable> -> !ttg.memdesc<64xf32, #shared1, #smem, mutable, 2x64>
%200 = tt.splat %183 : i1 -> tensor<64xi1, #blocked3>
%201 = arith.andi %200, %198 : tensor<64xi1, #blocked3>
%202 = ttg.async_copy_global_to_local %195, %199 mask %201 : tensor<64x!tt.ptr, #blocked3> -> <64xf32, #shared1, #smem, mutable, 2x64>
%203 = ttg.async_commit_group %202
%204 = ttg.memdesc_subview %111[%c1_i32, %c0_i32] : !ttg.memdesc<2x64xf32, #shared1, #smem, mutable> -> !ttg.memdesc<64xf32, #shared1, #smem, mutable, 2x64>
%205 = ttg.async_copy_global_to_local %195, %204 mask %201 : tensor<64x!tt.ptr, #blocked3> -> <64xf32, #shared1, #smem, mutable, 2x64>
%206 = ttg.async_commit_group %205
%207 = tt.expand_dims %191 {axis = 0 : i32} : tensor<64xi64, #ttg.slice<{dim = 0, parent = #blocked2}>> -> tensor<1x64xi64, #blocked2>
%208 = arith.muli %207, %cst_6 : tensor<1x64xi64, #blocked2>
%209 = tt.broadcast %208 : tensor<1x64xi64, #blocked2> -> tensor<256x64xi64, #blocked2>
%210 = arith.addi %70, %209 : tensor<256x64xi64, #blocked2>
%211 = tt.addptr %69, %210 : tensor<256x64x!tt.ptr, #blocked2>, tensor<256x64xi64, #blocked2>
%212 = arith.cmpi sge, %207, %cst_12 : tensor<1x64xi64, #blocked2>
%213 = arith.cmpi slt, %207, %71 : tensor<1x64xi64, #blocked2>
%214 = arith.andi %212, %213 : tensor<1x64xi1, #blocked2>
%215 = tt.broadcast %214 : tensor<1x64xi1, #blocked2> -> tensor<256x64xi1, #blocked2>
%216 = arith.andi %51, %215 : tensor<256x64xi1, #blocked2>
%217 = ttg.memdesc_subview %112[%c1_i32, %c0_i32, %c0_i32] : !ttg.memdesc<2x256x64xbf16, #shared2, #smem, mutable> -> !ttg.memdesc<256x64xbf16, #shared2, #smem, mutable, 2x256x64>
%218 = tt.splat %183 : i1 -> tensor<256x64xi1, #blocked2>
%219 = arith.andi %218, %216 : tensor<256x64xi1, #blocked2>
%220 = ttg.async_copy_global_to_local %211, %217 mask %219 : tensor<256x64x!tt.ptr, #blocked2> -> <256x64xbf16, #shared2, #smem, mutable, 2x256x64>
%221 = ttg.async_commit_group %220
%222 = tt.expand_dims %192 {axis = 1 : i32} : tensor<64xi64, #ttg.slice<{dim = 1, parent = #blocked}>> -> tensor<64x1xi64, #blocked>
%223 = tt.expand_dims %193 {axis = 1 : i32} : tensor<64xi64, #ttg.slice<{dim = 1, parent = #blocked1}>> -> tensor<64x1xi64, #blocked1>
%224 = arith.muli %222, %cst_4 : tensor<64x1xi64, #blocked>
%225 = arith.muli %223, %cst_5 : tensor<64x1xi64, #blocked1>
%226 = tt.broadcast %224 : tensor<64x1xi64, #blocked> -> tensor<64x256xi64, #blocked>
%227 = arith.addi %226, %74 : tensor<64x256xi64, #blocked>
%228 = tt.addptr %72, %227 : tensor<64x256x!tt.ptr, #blocked>, tensor<64x256xi64, #blocked>
%229 = arith.cmpi sge, %222, %cst_2 : tensor<64x1xi64, #blocked>
%230 = arith.cmpi sge, %223, %cst_3 : tensor<64x1xi64, #blocked1>
%231 = arith.cmpi slt, %222, %75 : tensor<64x1xi64, #blocked>
%232 = arith.cmpi slt, %223, %76 : tensor<64x1xi64, #blocked1>
%233 = arith.andi %229, %231 : tensor<64x1xi1, #blocked>
%234 = arith.andi %230, %232 : tensor<64x1xi1, #blocked1>
%235 = tt.broadcast %233 : tensor<64x1xi1, #blocked> -> tensor<64x256xi1, #blocked>
%236 = arith.andi %235, %80 : tensor<64x256xi1, #blocked>
%237 = ttg.memdesc_subview %113[%c1_i32, %c0_i32, %c0_i32] : !ttg.memdesc<2x64x256xbf16, #shared3, #smem, mutable> -> !ttg.memdesc<64x256xbf16, #shared3, #smem, mutable, 2x64x256>
%238 = tt.splat %183 : i1 -> tensor<64x256xi1, #blocked>
%239 = arith.andi %238, %236 : tensor<64x256xi1, #blocked>
%240 = ttg.async_copy_global_to_local %228, %237 mask %239 : tensor<64x256x!tt.ptr, #blocked> -> <64x256xbf16, #shared3, #smem, mutable, 2x64x256>
%241 = ttg.async_commit_group %240
%242 = tt.broadcast %225 : tensor<64x1xi64, #blocked1> -> tensor<64x64xi64, #blocked1>
%243 = arith.addi %242, %82 : tensor<64x64xi64, #blocked1>
%244 = tt.addptr %81, %243 : tensor<64x64x!tt.ptr, #blocked1>, tensor<64x64xi64, #blocked1>
%245 = tt.broadcast %234 : tensor<64x1xi1, #blocked1> -> tensor<64x64xi1, #blocked1>
%246 = arith.andi %245, %83 : tensor<64x64xi1, #blocked1>
%247 = ttg.memdesc_subview %114[%c1_i32, %c0_i32, %c0_i32] : !ttg.memdesc<2x64x64xbf16, #shared4, #smem, mutable> -> !ttg.memdesc<64x64xbf16, #shared4, #smem, mutable, 2x64x64>
%248 = tt.splat %183 : i1 -> tensor<64x64xi1, #blocked1>
%249 = arith.andi %248, %246 : tensor<64x64xi1, #blocked1>
%250 = ttg.async_copy_global_to_local %244, %247 mask %249 : tensor<64x64x!tt.ptr, #blocked1> -> <64x64xbf16, #shared4, #smem, mutable, 2x64x64>
%251 = ttg.async_commit_group %250
%252 = ttg.async_wait %137, %152, %172, %182 {num = 5 : i32}
ttng.fence_async_shared {bCluster = false}
%253:18 = scf.for %arg9 = %c0_i32 to %108 step %c1_i32 iter_args(%arg10 = %cst, %arg11 = %c1_i32, %arg12 = %c0_i32, %arg13 = %130, %arg14 = %252, %arg15 = %135, %arg16 = %148, %arg17 = %252, %arg18 = %168, %arg19 = %178, %arg20 = %174, %arg21 = %243, %arg22 = %177, %arg23 = %246, %arg24 = %206, %arg25 = %221, %arg26 = %241, %arg27 = %251) -> (tensor<256x64xf32, #mma>, i32, i32, !ttg.memdesc<64xf32, #shared1, #smem, mutable, 2x64>, !ttg.async.token, !ttg.memdesc<64xf32, #shared1, #smem, mutable, 2x64>, !ttg.memdesc<256x64xbf16, #shared2, #smem, mutable, 2x256x64>, !ttg.async.token, !ttg.memdesc<64x256xbf16, #shared3, #smem, mutable, 2x64x256>, !ttg.memdesc<64x64xbf16, #shared4, #smem, mutable, 2x64x64>, tensor<64x64xi64, #blocked1>, tensor<64x64xi64, #blocked1>, tensor<64x64xi1, #blocked1>, tensor<64x64xi1, #blocked1>, !ttg.async.token, !ttg.async.token, !ttg.async.token, !ttg.async.token) : i32 {
%259 = arith.subi %108, %c2_i32 : i32
%260 = arith.cmpi slt, %arg9, %259 : i32
%261 = ttg.local_load %arg13 token %arg14 : !ttg.memdesc<64xf32, #shared1, #smem, mutable, 2x64> -> tensor<64xf32, #ttg.slice<{dim = 0, parent = #blocked2}>>
%262 = ttg.local_load %arg15 : !ttg.memdesc<64xf32, #shared1, #smem, mutable, 2x64> -> tensor<64xf32, #ttg.slice<{dim = 1, parent = #ttg.dot_op<{opIdx = 0, parent = #mma1, kWidth = 2}>}>>
%263 = ttg.local_load %arg16 token %arg17 : !ttg.memdesc<256x64xbf16, #shared2, #smem, mutable, 2x256x64> -> tensor<256x64xbf16, #blocked2>
%264 = arith.subf %109, %261 : tensor<64xf32, #ttg.slice<{dim = 0, parent = #blocked2}>>
%265 = math.exp %264 : tensor<64xf32, #ttg.slice<{dim = 0, parent = #blocked2}>>
%266 = tt.expand_dims %265 {axis = 0 : i32} : tensor<64xf32, #ttg.slice<{dim = 0, parent = #blocked2}>> -> tensor<1x64xf32, #blocked2>
%267 = arith.extf %263 : tensor<256x64xbf16, #blocked2> to tensor<256x64xf32, #blocked2>
%268 = tt.broadcast %266 : tensor<1x64xf32, #blocked2> -> tensor<256x64xf32, #blocked2>
%269 = arith.mulf %267, %268 : tensor<256x64xf32, #blocked2>
%270 = arith.truncf %269 : tensor<256x64xf32, #blocked2> to tensor<256x64xbf16, #blocked2>
%271 = ttg.local_alloc %270 : (tensor<256x64xbf16, #blocked2>) -> !ttg.memdesc<256x64xbf16, #shared5, #smem>
%272 = ttg.local_load %arg18 : !ttg.memdesc<64x256xbf16, #shared3, #smem, mutable, 2x64x256> -> tensor<64x256xbf16, #ttg.dot_op<{opIdx = 0, parent = #mma1, kWidth = 2}>>
%273 = math.exp %262 : tensor<64xf32, #ttg.slice<{dim = 1, parent = #ttg.dot_op<{opIdx = 0, parent = #mma1, kWidth = 2}>}>>
%274 = tt.expand_dims %273 {axis = 1 : i32} : tensor<64xf32, #ttg.slice<{dim = 1, parent = #ttg.dot_op<{opIdx = 0, parent = #mma1, kWidth = 2}>}>> -> tensor<64x1xf32, #ttg.dot_op<{opIdx = 0, parent = #mma1, kWidth = 2}>>
%275 = arith.extf %272 : tensor<64x256xbf16, #ttg.dot_op<{opIdx = 0, parent = #mma1, kWidth = 2}>> to tensor<64x256xf32, #ttg.dot_op<{opIdx = 0, parent = #mma1, kWidth = 2}>>
%276 = tt.broadcast %274 : tensor<64x1xf32, #ttg.dot_op<{opIdx = 0, parent = #mma1, kWidth = 2}>> -> tensor<64x256xf32, #ttg.dot_op<{opIdx = 0, parent = #mma1, kWidth = 2}>>
%277 = arith.mulf %275, %276 : tensor<64x256xf32, #ttg.dot_op<{opIdx = 0, parent = #mma1, kWidth = 2}>>
%278 = arith.truncf %277 : tensor<64x256xf32, #ttg.dot_op<{opIdx = 0, parent = #mma1, kWidth = 2}>> to tensor<64x256xbf16, #ttg.dot_op<{opIdx = 0, parent = #mma1, kWidth = 2}>>
%279 = ttg.local_load %arg19 : !ttg.memdesc<64x64xbf16, #shared4, #smem, mutable, 2x64x64> -> tensor<64x64xbf16, #mma1>
%280 = ttng.warp_group_dot %278, %92, %cst_7, %false {inputPrecision = 0 : i32, isAsync = true} : tensor<64x256xbf16, #ttg.dot_op<{opIdx = 0, parent = #mma1, kWidth = 2}>> * !ttg.memdesc<256x64xbf16, #shared, #smem> -> tensor<64x64xf32, #mma1>
%281:2 = ttng.warp_group_dot_wait %280, %92 {pendings = 0 : i32} : tensor<64x64xf32, #mma1>, !ttg.memdesc<256x64xbf16, #shared, #smem>
%282 = arith.extf %279 : tensor<64x64xbf16, #mma1> to tensor<64x64xf32, #mma1>
%283 = arith.subf %282, %281#0 : tensor<64x64xf32, #mma1>
%284 = arith.truncf %283 : tensor<64x64xf32, #mma1> to tensor<64x64xbf16, #mma1>
%285 = ttg.local_alloc %284 : (tensor<64x64xbf16, #mma1>) -> !ttg.memdesc<64x64xbf16, #shared, #smem>
%286 = tt.addptr %84, %arg20 : tensor<64x64x!tt.ptr, #blocked1>, tensor<64x64xi64, #blocked1>
%287 = ttg.convert_layout %284 : tensor<64x64xbf16, #mma1> -> tensor<64x64xbf16, #blocked1>
tt.store %286, %287, %arg22 : tensor<64x64x!tt.ptr, #blocked1>
%288 = ttg.local_load %271 : !ttg.memdesc<256x64xbf16, #shared5, #smem> -> tensor<256x64xbf16, #ttg.dot_op<{opIdx = 0, parent = #mma, kWidth = 2}>>
ttng.fence_async_shared {bCluster = false}
%289 = ttng.warp_group_dot %288, %285, %arg10 {isAsync = true} : tensor<256x64xbf16, #ttg.dot_op<{opIdx = 0, parent = #mma, kWidth = 2}>> * !ttg.memdesc<64x64xbf16, #shared, #smem> -> tensor<256x64xf32, #mma>
%290:2 = ttng.warp_group_dot_wait %289, %285 {pendings = 0 : i32} : tensor<256x64xf32, #mma>, !ttg.memdesc<64x64xbf16, #shared, #smem>
%291 = arith.addi %arg11, %c1_i32 : i32
%292 = arith.cmpi slt, %291, %c2_i32 : i32
%293 = arith.select %292, %291, %c0_i32 : i32
%294 = arith.addi %arg9, %c2_i32 : i32
%295 = arith.muli %294, %c64_i32 : i32
%296 = arith.addi %104, %295 : i32
%297 = arith.extsi %296 : i32 to i64
%298 = tt.splat %297 : i64 -> tensor<64xi64, #blocked3>
%299 = tt.splat %297 : i64 -> tensor<64xi64, #ttg.slice<{dim = 0, parent = #blocked2}>>
%300 = tt.splat %297 : i64 -> tensor<64xi64, #ttg.slice<{dim = 1, parent = #blocked}>>
%301 = tt.splat %297 : i64 -> tensor<64xi64, #ttg.slice<{dim = 1, parent = #blocked1}>>
%302 = arith.addi %298, %36 : tensor<64xi64, #blocked3>
%303 = arith.addi %299, %37 : tensor<64xi64, #ttg.slice<{dim = 0, parent = #blocked2}>>
%304 = arith.addi %300, %38 : tensor<64xi64, #ttg.slice<{dim = 1, parent = #blocked}>>
%305 = arith.addi %301, %39 : tensor<64xi64, #ttg.slice<{dim = 1, parent = #blocked1}>>
%306 = arith.muli %302, %cst_15 : tensor<64xi64, #blocked3>
%307 = tt.addptr %67, %306 : tensor<64x!tt.ptr, #blocked3>, tensor<64xi64, #blocked3>
%308 = arith.cmpi sge, %302, %cst_16 : tensor<64xi64, #blocked3>
%309 = arith.cmpi slt, %302, %68 : tensor<64xi64, #blocked3>
%310 = arith.andi %308, %309 : tensor<64xi1, #blocked3>
%311 = ttg.memdesc_subview %110[%293, %c0_i32] : !ttg.memdesc<2x64xf32, #shared1, #smem, mutable> -> !ttg.memdesc<64xf32, #shared1, #smem, mutable, 2x64>
%312 = tt.splat %260 : i1 -> tensor<64xi1, #blocked3>
%313 = arith.andi %312, %310 : tensor<64xi1, #blocked3>
%314 = ttg.async_copy_global_to_local %307, %311 mask %313 : tensor<64x!tt.ptr, #blocked3> -> <64xf32, #shared1, #smem, mutable, 2x64>
%315 = ttg.async_commit_group %314
%316 = ttg.memdesc_subview %111[%293, %c0_i32] : !ttg.memdesc<2x64xf32, #shared1, #smem, mutable> -> !ttg.memdesc<64xf32, #shared1, #smem, mutable, 2x64>
%317 = ttg.async_copy_global_to_local %307, %316 mask %313 : tensor<64x!tt.ptr, #blocked3> -> <64xf32, #shared1, #smem, mutable, 2x64>
%318 = ttg.async_commit_group %317
%319 = tt.expand_dims %303 {axis = 0 : i32} : tensor<64xi64, #ttg.slice<{dim = 0, parent = #blocked2}>> -> tensor<1x64xi64, #blocked2>
%320 = arith.muli %319, %cst_6 : tensor<1x64xi64, #blocked2>
%321 = tt.broadcast %320 : tensor<1x64xi64, #blocked2> -> tensor<256x64xi64, #blocked2>
%322 = arith.addi %70, %321 : tensor<256x64xi64, #blocked2>
%323 = tt.addptr %69, %322 : tensor<256x64x!tt.ptr, #blocked2>, tensor<256x64xi64, #blocked2>
%324 = arith.cmpi sge, %319, %cst_12 : tensor<1x64xi64, #blocked2>
%325 = arith.cmpi slt, %319, %71 : tensor<1x64xi64, #blocked2>
%326 = arith.andi %324, %325 : tensor<1x64xi1, #blocked2>
%327 = tt.broadcast %326 : tensor<1x64xi1, #blocked2> -> tensor<256x64xi1, #blocked2>
%328 = arith.andi %51, %327 : tensor<256x64xi1, #blocked2>
%329 = ttg.memdesc_subview %112[%293, %c0_i32, %c0_i32] : !ttg.memdesc<2x256x64xbf16, #shared2, #smem, mutable> -> !ttg.memdesc<256x64xbf16, #shared2, #smem, mutable, 2x256x64>
%330 = tt.splat %260 : i1 -> tensor<256x64xi1, #blocked2>
%331 = arith.andi %330, %328 : tensor<256x64xi1, #blocked2>
%332 = ttg.async_copy_global_to_local %323, %329 mask %331 : tensor<256x64x!tt.ptr, #blocked2> -> <256x64xbf16, #shared2, #smem, mutable, 2x256x64>
%333 = ttg.async_commit_group %332
%334 = tt.expand_dims %304 {axis = 1 : i32} : tensor<64xi64, #ttg.slice<{dim = 1, parent = #blocked}>> -> tensor<64x1xi64, #blocked>
%335 = tt.expand_dims %305 {axis = 1 : i32} : tensor<64xi64, #ttg.slice<{dim = 1, parent = #blocked1}>> -> tensor<64x1xi64, #blocked1>
%336 = arith.muli %334, %cst_4 : tensor<64x1xi64, #blocked>
%337 = arith.muli %335, %cst_5 : tensor<64x1xi64, #blocked1>
%338 = tt.broadcast %336 : tensor<64x1xi64, #blocked> -> tensor<64x256xi64, #blocked>
%339 = arith.addi %338, %74 : tensor<64x256xi64, #blocked>
%340 = tt.addptr %72, %339 : tensor<64x256x!tt.ptr, #blocked>, tensor<64x256xi64, #blocked>
%341 = arith.cmpi sge, %334, %cst_2 : tensor<64x1xi64, #blocked>
%342 = arith.cmpi sge, %335, %cst_3 : tensor<64x1xi64, #blocked1>
%343 = arith.cmpi slt, %334, %75 : tensor<64x1xi64, #blocked>
%344 = arith.cmpi slt, %335, %76 : tensor<64x1xi64, #blocked1>
%345 = arith.andi %341, %343 : tensor<64x1xi1, #blocked>
%346 = arith.andi %342, %344 : tensor<64x1xi1, #blocked1>
%347 = tt.broadcast %345 : tensor<64x1xi1, #blocked> -> tensor<64x256xi1, #blocked>
%348 = arith.andi %347, %80 : tensor<64x256xi1, #blocked>
%349 = ttg.memdesc_subview %113[%293, %c0_i32, %c0_i32] : !ttg.memdesc<2x64x256xbf16, #shared3, #smem, mutable> -> !ttg.memdesc<64x256xbf16, #shared3, #smem, mutable, 2x64x256>
%350 = tt.splat %260 : i1 -> tensor<64x256xi1, #blocked>
%351 = arith.andi %350, %348 : tensor<64x256xi1, #blocked>
%352 = ttg.async_copy_global_to_local %340, %349 mask %351 : tensor<64x256x!tt.ptr, #blocked> -> <64x256xbf16, #shared3, #smem, mutable, 2x64x256>
%353 = ttg.async_commit_group %352
%354 = tt.broadcast %337 : tensor<64x1xi64, #blocked1> -> tensor<64x64xi64, #blocked1>
%355 = arith.addi %354, %82 : tensor<64x64xi64, #blocked1>
%356 = tt.addptr %81, %355 : tensor<64x64x!tt.ptr, #blocked1>, tensor<64x64xi64, #blocked1>
%357 = tt.broadcast %346 : tensor<64x1xi1, #blocked1> -> tensor<64x64xi1, #blocked1>
%358 = arith.andi %357, %83 : tensor<64x64xi1, #blocked1>
%359 = ttg.memdesc_subview %114[%293, %c0_i32, %c0_i32] : !ttg.memdesc<2x64x64xbf16, #shared4, #smem, mutable> -> !ttg.memdesc<64x64xbf16, #shared4, #smem, mutable, 2x64x64>
%360 = tt.splat %260 : i1 -> tensor<64x64xi1, #blocked1>
%361 = arith.andi %360, %358 : tensor<64x64xi1, #blocked1>
%362 = ttg.async_copy_global_to_local %356, %359 mask %361 : tensor<64x64x!tt.ptr, #blocked1> -> <64x64xbf16, #shared4, #smem, mutable, 2x64x64>
%363 = ttg.async_commit_group %362
%364 = arith.addi %arg12, %c1_i32 : i32
%365 = arith.cmpi slt, %364, %c2_i32 : i32
%366 = arith.select %365, %364, %c0_i32 : i32
%367 = ttg.memdesc_subview %110[%366, %c0_i32] : !ttg.memdesc<2x64xf32, #shared1, #smem, mutable> -> !ttg.memdesc<64xf32, #shared1, #smem, mutable, 2x64>
%368 = ttg.memdesc_subview %111[%366, %c0_i32] : !ttg.memdesc<2x64xf32, #shared1, #smem, mutable> -> !ttg.memdesc<64xf32, #shared1, #smem, mutable, 2x64>
%369 = ttg.memdesc_subview %112[%366, %c0_i32, %c0_i32] : !ttg.memdesc<2x256x64xbf16, #shared2, #smem, mutable> -> !ttg.memdesc<256x64xbf16, #shared2, #smem, mutable, 2x256x64>
%370 = ttg.memdesc_subview %113[%366, %c0_i32, %c0_i32] : !ttg.memdesc<2x64x256xbf16, #shared3, #smem, mutable> -> !ttg.memdesc<64x256xbf16, #shared3, #smem, mutable, 2x64x256>
%371 = ttg.async_wait %arg24, %arg25, %arg26, %arg27 {num = 5 : i32}
%372 = ttg.memdesc_subview %114[%366, %c0_i32, %c0_i32] : !ttg.memdesc<2x64x64xbf16, #shared4, #smem, mutable> -> !ttg.memdesc<64x64xbf16, #shared4, #smem, mutable, 2x64x64>
scf.yield %290#0, %293, %366, %367, %371, %368, %369, %371, %370, %372, %arg21, %355, %arg23, %358, %318, %333, %353, %363 : tensor<256x64xf32, #mma>, i32, i32, !ttg.memdesc<64xf32, #shared1, #smem, mutable, 2x64>, !ttg.async.token, !ttg.memdesc<64xf32, #shared1, #smem, mutable, 2x64>, !ttg.memdesc<256x64xbf16, #shared2, #smem, mutable, 2x256x64>, !ttg.async.token, !ttg.memdesc<64x256xbf16, #shared3, #smem, mutable, 2x64x256>, !ttg.memdesc<64x64xbf16, #shared4, #smem, mutable, 2x64x64>, tensor<64x64xi64, #blocked1>, tensor<64x64xi64, #blocked1>, tensor<64x64xi1, #blocked1>, tensor<64x64xi1, #blocked1>, !ttg.async.token, !ttg.async.token, !ttg.async.token, !ttg.async.token
}
%254 = ttg.async_wait {num = 0 : i32}
ttg.local_dealloc %110 : !ttg.memdesc<2x64xf32, #shared1, #smem, mutable>
ttg.local_dealloc %111 : !ttg.memdesc<2x64xf32, #shared1, #smem, mutable>
ttg.local_dealloc %112 : !ttg.memdesc<2x256x64xbf16, #shared2, #smem, mutable>
ttg.local_dealloc %113 : !ttg.memdesc<2x64x256xbf16, #shared3, #smem, mutable>
ttg.local_dealloc %114 : !ttg.memdesc<2x64x64xbf16, #shared4, #smem, mutable>
%255 = math.exp %103 : f32
%256 = tt.splat %255 : f32 -> tensor<256x64xf32, #mma>
%257 = arith.mulf %arg8, %256 : tensor<256x64xf32, #mma>
%258 = arith.addf %257, %253#0 : tensor<256x64xf32, #mma>
scf.yield %258 : tensor<256x64xf32, #mma>
}
tt.return
}
}

{-#
external_resources: {
mlir_reproducer: {
pipeline: "builtin.module(triton-nvidia-mma-lowering, tritongpu-combine-tensor-select-and-if, convert-scf-to-cf, allocate-shared-memory, triton-tensor-memory-allocation, tritongpu-global-scratch-memory-allocation, convert-triton-gpu-to-llvm{compute-capability=90 ptx-version=84}, canonicalize{ max-iterations=10 max-num-rewrites=-1 region-simplify=normal test-convergence=false top-down=true}, cse, convert-nv-gpu-to-llvm, canonicalize{ max-iterations=10 max-num-rewrites=-1 region-simplify=normal test-convergence=false top-down=true}, cse, symbol-dce, enable-line-info)",
disable_threading: false,
verify_each: true
}
}
#-}
/usr/local/lib/python3.10/dist-packages/fla/ops/common/chunk_delta_h.py:27:0: error: Failures have been detected while processing an MLIR pass pipeline
/usr/local/lib/python3.10/dist-packages/fla/ops/common/chunk_delta_h.py:27:0: note: Pipeline failed while executing [ConvertTritonGPUToLLVM on 'builtin.module' operation]: reproducer generated at std::errs, please share the reproducer above with Triton project.
Traceback (most recent call last):
File "/data/Megatron-LM/reproduce.py", line 31, in
chunk_gated_delta_rule_fwd_h(k=k, w=w, u=u, g=g, head_first=head_first, chunk_size=chunk_size)
File "/usr/local/lib/python3.10/dist-packages/fla/ops/common/chunk_delta_h.py", line 290, in chunk_gated_delta_rule_fwd_h
chunk_gated_delta_rule_fwd_kernel_h[grid](
File "/workspace/triton/python/triton/runtime/jit.py", line 336, in
return lambda *args, **kwargs: self.run(grid=grid, warmup=False, *args, **kwargs)
File "/workspace/triton/python/triton/runtime/autotuner.py", line 379, in run
return self.fn.run(*args, **kwargs)
File "/workspace/triton/python/triton/runtime/autotuner.py", line 189, in run
timings = {config: self._bench(*args, config=config, **kwargs) for config in pruned_configs}
File "/workspace/triton/python/triton/runtime/autotuner.py", line 189, in
timings = {config: self._bench(*args, config=config, **kwargs) for config in pruned_configs}
File "/workspace/triton/python/triton/runtime/autotuner.py", line 167, in _bench
return self.do_bench(kernel_call, quantiles=(0.5, 0.2, 0.8))
File "/workspace/triton/python/triton/testing.py", line 145, in do_bench
fn()
File "/workspace/triton/python/triton/runtime/autotuner.py", line 153, in kernel_call
self.fn.run(
File "/workspace/triton/python/triton/runtime/jit.py", line 563, in run
kernel = self.compile(src, target=target, options=options.dict)
File "/workspace/triton/python/triton/compiler/compiler.py", line 283, in compile
next_module = compile_ir(module, metadata)
File "/workspace/triton/python/triton/backends/nvidia/compiler.py", line 431, in
stages["llir"] = lambda src, metadata: self.make_llir(src, metadata, options, capability)
File "/workspace/triton/python/triton/backends/nvidia/compiler.py", line 322, in make_llir
pm.run(mod)
RuntimeError: PassManager::run failed

[sustcsonglin]

oh i mean reinstall from the source https://github.com/triton-lang/triton?tab=readme-ov-file#install-from-source

removing num_warps = 8 can workaround, but the speed could be greatly compromised.. I'd suggest reduce the block size if one has to use num_warps=4 for this Triton compiler bug

[LouChao98]

yeah, I compiled triton from source following the instructions exactly. I dont know how "from scratch" comes into my mind... It is a typo.

[sustcsonglin]

Well I am not sure what do they mean by already fixed in the main branch triton-lang/triton#5609 (comment)

[sustcsonglin]

I assume this issue has been fixed using some nightly version of Triton. Feel free to reopen it when necessary