#blocked0 = #triton_gpu.blocked<{sizePerThread = [1], threadsPerWarp = [32], warpsPerCTA = [8], order = [0]}>
#blocked1 = #triton_gpu.blocked<{sizePerThread = [1, 8], threadsPerWarp = [4, 8], warpsPerCTA = [8, 1], order = [1, 0]}>
#blocked2 = #triton_gpu.blocked<{sizePerThread = [1, 4], threadsPerWarp = [2, 16], warpsPerCTA = [8, 1], order = [1, 0]}>
#blocked3 = #triton_gpu.blocked<{sizePerThread = [1, 1], threadsPerWarp = [32, 1], warpsPerCTA = [4, 2], order = [0, 1]}>
#mma0 = #triton_gpu.mma<{versionMajor = 2, versionMinor = 0, warpsPerCTA = [4, 2]}>
#mma1 = #triton_gpu.mma<{versionMajor = 2, versionMinor = 0, warpsPerCTA = [8, 1]}>
#shared0 = #triton_gpu.shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [1, 0]}>
#shared1 = #triton_gpu.shared<{vec = 1, perPhase = 1, maxPhase = 1, order = [0, 1]}>
module attributes {"triton_gpu.num-warps" = 8 : i32} {
  func public @_bwd_kernel_0d1d2d34d5d6d7d8d9d10d11d12d13d14d15c16d17d18d19c20d21d22d23c2425d26d27(%arg0: !tt.ptr<f16> {tt.divisibility = 16 : i32}, %arg1: !tt.ptr<f16> {tt.divisibility = 16 : i32}, %arg2: !tt.ptr<f16> {tt.divisibility = 16 : i32}, %arg3: f32, %arg4: !tt.ptr<f16> {tt.divisibility = 16 : i32}, %arg5: !tt.ptr<f16> {tt.divisibility = 16 : i32}, %arg6: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg7: !tt.ptr<f16> {tt.divisibility = 16 : i32}, %arg8: !tt.ptr<f16> {tt.divisibility = 16 : i32}, %arg9: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg10: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg11: !tt.ptr<f32> {tt.divisibility = 16 : i32}, %arg12: i32 {tt.divisibility = 16 : i32}, %arg13: i32 {tt.divisibility = 16 : i32}, %arg14: i32 {tt.divisibility = 16 : i32}, %arg15: i32 {tt.divisibility = 16 : i32}, %arg16: i32 {tt.divisibility = 16 : i32}, %arg17: i32 {tt.divisibility = 16 : i32}, %arg18: i32 {tt.divisibility = 16 : i32}, %arg19: i32 {tt.divisibility = 16 : i32}, %arg20: i32 {tt.divisibility = 16 : i32}, %arg21: i32, %arg22: i32 {tt.divisibility = 16 : i32}, %arg23: i32 {tt.divisibility = 16 : i32}, %arg24: i32) {
    %c0 = arith.constant 0 : index
    %c1 = arith.constant 1 : index
    %c128_i32 = arith.constant 128 : i32
    %c128 = arith.constant 128 : index
    %cst = arith.constant dense<0xFF800000> : tensor<128x128xf32, #blocked3>
    %cst_0 = arith.constant dense<0.000000e+00> : tensor<128x64xf32, #mma0>
    %cst_1 = arith.constant dense<0.000000e+00> : tensor<128x128xf32, #mma1>
    %0 = tt.get_program_id {axis = 0 : i32} : i32
    %1 = arith.divsi %0, %arg22 : i32
    %2 = arith.remsi %0, %arg22 : i32
    %3 = arith.muli %1, %arg12 : i32
    %4 = arith.muli %2, %arg13 : i32
    %5 = arith.addi %3, %4 : i32
    %6 = tt.addptr %arg0, %5 : !tt.ptr<f16>, i32
    %7 = tt.addptr %arg1, %5 : !tt.ptr<f16>, i32
    %8 = tt.addptr %arg2, %5 : !tt.ptr<f16>, i32
    %9 = tt.addptr %arg5, %5 : !tt.ptr<f16>, i32
    %10 = tt.addptr %arg6, %5 : !tt.ptr<f32>, i32
    %11 = tt.addptr %arg7, %5 : !tt.ptr<f16>, i32
    %12 = tt.addptr %arg8, %5 : !tt.ptr<f16>, i32
    %13 = arith.index_cast %arg24 : i32 to index
    %14 = tt.make_range {end = 128 : i32, start = 0 : i32} : tensor<128xi32, #blocked0>
    %15 = tt.make_range {end = 128 : i32, start = 0 : i32} : tensor<128xi32, #triton_gpu.slice<{dim = 1, parent = #blocked1}>>
    %16 = tt.make_range {end = 128 : i32, start = 0 : i32} : tensor<128xi32, #triton_gpu.slice<{dim = 0, parent = #blocked3}>>
    %17 = tt.make_range {end = 128 : i32, start = 0 : i32} : tensor<128xi32, #triton_gpu.slice<{dim = 1, parent = #blocked2}>>
    %18 = tt.make_range {end = 128 : i32, start = 0 : i32} : tensor<128xi32, #triton_gpu.slice<{dim = 1, parent = #blocked3}>>
    %19 = tt.splat %arg14 : (i32) -> tensor<128x1xi32, #blocked1>
    %20 = tt.splat %arg14 : (i32) -> tensor<128x1xi32, #blocked2>
    %21 = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32, #triton_gpu.slice<{dim = 0, parent = #blocked1}>>
    %22 = tt.make_range {end = 64 : i32, start = 0 : i32} : tensor<64xi32, #triton_gpu.slice<{dim = 0, parent = #blocked2}>>
    %23 = tt.expand_dims %21 {axis = 0 : i32} : (tensor<64xi32, #triton_gpu.slice<{dim = 0, parent = #blocked1}>>) -> tensor<1x64xi32, #blocked1>
    %24 = tt.broadcast %23 : (tensor<1x64xi32, #blocked1>) -> tensor<128x64xi32, #blocked1>
    %25 = tt.expand_dims %22 {axis = 0 : i32} : (tensor<64xi32, #triton_gpu.slice<{dim = 0, parent = #blocked2}>>) -> tensor<1x64xi32, #blocked2>
    %26 = tt.broadcast %25 : (tensor<1x64xi32, #blocked2>) -> tensor<128x64xi32, #blocked2>
    %27 = tt.splat %6 : (!tt.ptr<f16>) -> tensor<128x64x!tt.ptr<f16>, #blocked1>
    %28 = tt.splat %arg17 : (i32) -> tensor<128x1xi32, #blocked1>
    %29 = tt.splat %7 : (!tt.ptr<f16>) -> tensor<128x64x!tt.ptr<f16>, #blocked1>
    %30 = tt.splat %8 : (!tt.ptr<f16>) -> tensor<128x64x!tt.ptr<f16>, #blocked1>
    %31 = tt.splat %9 : (!tt.ptr<f16>) -> tensor<128x64x!tt.ptr<f16>, #blocked1>
    %32 = tt.splat %10 : (!tt.ptr<f32>) -> tensor<128x64x!tt.ptr<f32>, #blocked2>
    %33 = arith.muli %0, %arg23 : i32
    %34 = tt.addptr %arg11, %33 : !tt.ptr<f32>, i32
    %35 = tt.addptr %arg10, %33 : !tt.ptr<f32>, i32
    %36 = arith.muli %arg24, %c128_i32 : i32
    %37 = arith.index_cast %36 : i32 to index
    %38 = tt.splat %35 : (!tt.ptr<f32>) -> tensor<128x!tt.ptr<f32>, #blocked0>
    %39 = tt.splat %arg3 : (f32) -> tensor<128x128xf32, #blocked3>
    %40 = tt.splat %34 : (!tt.ptr<f32>) -> tensor<128x!tt.ptr<f32>, #blocked0>
    %41 = arith.muli %arg14, %c128_i32 : i32
    %42 = tt.splat %41 : (i32) -> tensor<128x64xi32, #blocked1>
    %43 = tt.splat %41 : (i32) -> tensor<128x64xi32, #blocked2>
    %44 = tt.splat %12 : (!tt.ptr<f16>) -> tensor<128x64x!tt.ptr<f16>, #blocked1>
    %45 = tt.splat %11 : (!tt.ptr<f16>) -> tensor<128x64x!tt.ptr<f16>, #blocked1>
    scf.for %arg25 = %c0 to %13 step %c1 {
      %46 = arith.index_cast %arg25 : index to i32
      %47 = arith.muli %46, %c128_i32 : i32
      %48 = tt.splat %47 : (i32) -> tensor<128xi32, #triton_gpu.slice<{dim = 1, parent = #blocked1}>>
      %49 = tt.splat %47 : (i32) -> tensor<128xi32, #triton_gpu.slice<{dim = 0, parent = #blocked3}>>
      %50 = tt.splat %47 : (i32) -> tensor<128xi32, #triton_gpu.slice<{dim = 1, parent = #blocked2}>>
      %51 = arith.addi %48, %15 : tensor<128xi32, #triton_gpu.slice<{dim = 1, parent = #blocked1}>>
      %52 = arith.addi %50, %17 : tensor<128xi32, #triton_gpu.slice<{dim = 1, parent = #blocked2}>>
      %53 = tt.expand_dims %51 {axis = 1 : i32} : (tensor<128xi32, #triton_gpu.slice<{dim = 1, parent = #blocked1}>>) -> tensor<128x1xi32, #blocked1>
      %54 = tt.expand_dims %52 {axis = 1 : i32} : (tensor<128xi32, #triton_gpu.slice<{dim = 1, parent = #blocked2}>>) -> tensor<128x1xi32, #blocked2>
      %55 = arith.muli %53, %28 : tensor<128x1xi32, #blocked1>
      %56 = tt.broadcast %55 : (tensor<128x1xi32, #blocked1>) -> tensor<128x64xi32, #blocked1>
      %57 = arith.addi %56, %24 : tensor<128x64xi32, #blocked1>
      %58 = tt.addptr %29, %57 : tensor<128x64x!tt.ptr<f16>, #blocked1>, tensor<128x64xi32, #blocked1>
      %59 = tt.load %58 {cache = 1 : i32, evict = 1 : i32, isVolatile = false} : tensor<128x64xf16, #blocked1>
      %60 = arith.muli %53, %19 : tensor<128x1xi32, #blocked1>
      %61 = tt.broadcast %60 : (tensor<128x1xi32, #blocked1>) -> tensor<128x64xi32, #blocked1>
      %62 = arith.addi %61, %24 : tensor<128x64xi32, #blocked1>
      %63 = tt.addptr %30, %62 : tensor<128x64x!tt.ptr<f16>, #blocked1>, tensor<128x64xi32, #blocked1>
      %64 = tt.load %63 {cache = 1 : i32, evict = 1 : i32, isVolatile = false} : tensor<128x64xf16, #blocked1>
      %65 = arith.index_cast %47 : i32 to index
      %66 = triton_gpu.convert_layout %59 : (tensor<128x64xf16, #blocked1>) -> tensor<128x64xf16, #shared0>
      %67 = tt.trans %66 : (tensor<128x64xf16, #shared0>) -> tensor<64x128xf16, #shared1>
      %68 = arith.addi %49, %16 : tensor<128xi32, #triton_gpu.slice<{dim = 0, parent = #blocked3}>>
      %69 = tt.expand_dims %68 {axis = 0 : i32} : (tensor<128xi32, #triton_gpu.slice<{dim = 0, parent = #blocked3}>>) -> tensor<1x128xi32, #blocked3>
      %70 = tt.broadcast %69 : (tensor<1x128xi32, #blocked3>) -> tensor<128x128xi32, #blocked3>
      %71 = triton_gpu.convert_layout %64 : (tensor<128x64xf16, #blocked1>) -> tensor<128x64xf16, #shared0>
      %72 = tt.trans %71 : (tensor<128x64xf16, #shared0>) -> tensor<64x128xf16, #shared1>
      %73 = arith.muli %54, %20 : tensor<128x1xi32, #blocked2>
      %74 = tt.broadcast %73 : (tensor<128x1xi32, #blocked2>) -> tensor<128x64xi32, #blocked2>
      %75 = arith.addi %74, %26 : tensor<128x64xi32, #blocked2>
      %76 = tt.addptr %32, %75 : tensor<128x64x!tt.ptr<f32>, #blocked2>, tensor<128x64xi32, #blocked2>
      %77 = tt.addptr %27, %62 : tensor<128x64x!tt.ptr<f16>, #blocked1>, tensor<128x64xi32, #blocked1>
      %78 = tt.addptr %31, %62 : tensor<128x64x!tt.ptr<f16>, #blocked1>, tensor<128x64xi32, #blocked1>
      %79 = triton_gpu.convert_layout %67 : (tensor<64x128xf16, #shared1>) -> tensor<64x128xf16, #triton_gpu.dot_op<{opIdx = 1, parent = #mma1}>>
      %82:5 = scf.for %arg26 = %65 to %37 step %c128 iter_args(%arg27 = %cst_0, %arg28 = %cst_0, %arg29 = %76, %arg30 = %77, %arg31 = %78) -> (tensor<128x64xf32, #mma0>, tensor<128x64xf32, #mma0>, tensor<128x64x!tt.ptr<f32>, #blocked2>, tensor<128x64x!tt.ptr<f16>, #blocked1>, tensor<128x64x!tt.ptr<f16>, #blocked1>) {
        %89 = arith.index_cast %arg26 : index to i32
        %90 = tt.splat %89 : (i32) -> tensor<128xi32, #blocked0>
        %91 = tt.splat %89 : (i32) -> tensor<128xi32, #triton_gpu.slice<{dim = 1, parent = #blocked3}>>
        %92 = arith.addi %90, %14 : tensor<128xi32, #blocked0>
        %93 = tt.load %arg30 {cache = 1 : i32, evict = 1 : i32, isVolatile = false} : tensor<128x64xf16, #blocked1>
        %94 = triton_gpu.convert_layout %93 : (tensor<128x64xf16, #blocked1>) -> tensor<128x64xf16, #triton_gpu.dot_op<{opIdx = 0, parent = #mma1}>>
        %95 = tt.dot %94, %79, %cst_1 {allowTF32 = true} : tensor<128x64xf16, #triton_gpu.dot_op<{opIdx = 0, parent = #mma1}>> * tensor<64x128xf16, #triton_gpu.dot_op<{opIdx = 1, parent = #mma1}>> -> tensor<128x128xf32, #mma1>
        %96 = triton_gpu.convert_layout %95 : (tensor<128x128xf32, #mma1>) -> tensor<128x128xf32, #blocked3>
        %97 = arith.addi %91, %18 : tensor<128xi32, #triton_gpu.slice<{dim = 1, parent = #blocked3}>>
        %98 = tt.expand_dims %97 {axis = 1 : i32} : (tensor<128xi32, #triton_gpu.slice<{dim = 1, parent = #blocked3}>>) -> tensor<128x1xi32, #blocked3>
        %99 = tt.broadcast %98 : (tensor<128x1xi32, #blocked3>) -> tensor<128x128xi32, #blocked3>
        %100 = "triton_gpu.cmpi"(%99, %70) {predicate = 5 : i64} : (tensor<128x128xi32, #blocked3>, tensor<128x128xi32, #blocked3>) -> tensor<128x128xi1, #blocked3>
        %101 = "triton_gpu.select"(%100, %96, %cst) : (tensor<128x128xi1, #blocked3>, tensor<128x128xf32, #blocked3>, tensor<128x128xf32, #blocked3>) -> tensor<128x128xf32, #blocked3>
        %102 = tt.addptr %38, %92 : tensor<128x!tt.ptr<f32>, #blocked0>, tensor<128xi32, #blocked0>
        %103 = tt.load %102 {cache = 1 : i32, evict = 1 : i32, isVolatile = false} : tensor<128xf32, #blocked0>
        %104 = arith.mulf %101, %39 : tensor<128x128xf32, #blocked3>
        %105 = triton_gpu.convert_layout %103 : (tensor<128xf32, #blocked0>) -> tensor<128xf32, #triton_gpu.slice<{dim = 1, parent = #blocked3}>>
        %106 = tt.expand_dims %105 {axis = 1 : i32} : (tensor<128xf32, #triton_gpu.slice<{dim = 1, parent = #blocked3}>>) -> tensor<128x1xf32, #blocked3>
        %107 = tt.broadcast %106 : (tensor<128x1xf32, #blocked3>) -> tensor<128x128xf32, #blocked3>
        %108 = arith.subf %104, %107 : tensor<128x128xf32, #blocked3>
        %109 = math.exp %108 : tensor<128x128xf32, #blocked3>
        %110 = tt.load %arg31 {cache = 1 : i32, evict = 1 : i32, isVolatile = false} : tensor<128x64xf16, #blocked1>
        %111 = arith.truncf %109 : tensor<128x128xf32, #blocked3> to tensor<128x128xf16, #blocked3>
        %112 = triton_gpu.convert_layout %111 : (tensor<128x128xf16, #blocked3>) -> tensor<128x128xf16, #shared1>
        %113 = tt.trans %112 : (tensor<128x128xf16, #shared1>) -> tensor<128x128xf16, #shared0>
        %114 = triton_gpu.convert_layout %113 : (tensor<128x128xf16, #shared0>) -> tensor<128x128xf16, #triton_gpu.dot_op<{opIdx = 0, parent = #mma0}>>
        %115 = triton_gpu.convert_layout %110 : (tensor<128x64xf16, #blocked1>) -> tensor<128x64xf16, #triton_gpu.dot_op<{opIdx = 1, parent = #mma0}>>
        %116 = tt.dot %114, %115, %arg27 {allowTF32 = true} : tensor<128x128xf16, #triton_gpu.dot_op<{opIdx = 0, parent = #mma0}>> * tensor<128x64xf16, #triton_gpu.dot_op<{opIdx = 1, parent = #mma0}>> -> tensor<128x64xf32, #mma0>
        %117 = tt.addptr %40, %92 : tensor<128x!tt.ptr<f32>, #blocked0>, tensor<128xi32, #blocked0>
        %118 = tt.load %117 {cache = 1 : i32, evict = 1 : i32, isVolatile = false} : tensor<128xf32, #blocked0>
        %119 = triton_gpu.convert_layout %118 : (tensor<128xf32, #blocked0>) -> tensor<128xf32, #triton_gpu.slice<{dim = 1, parent = #mma1}>>
        %120 = tt.expand_dims %119 {axis = 1 : i32} : (tensor<128xf32, #triton_gpu.slice<{dim = 1, parent = #mma1}>>) -> tensor<128x1xf32, #mma1>
        %121 = tt.broadcast %120 : (tensor<128x1xf32, #mma1>) -> tensor<128x128xf32, #mma1>
        %122 = arith.subf %cst_1, %121 : tensor<128x128xf32, #mma1>
        %123 = triton_gpu.convert_layout %110 : (tensor<128x64xf16, #blocked1>) -> tensor<128x64xf16, #triton_gpu.dot_op<{opIdx = 0, parent = #mma1}>>
        %80 = triton_gpu.convert_layout %72 : (tensor<64x128xf16, #shared1>) -> tensor<64x128xf16, #triton_gpu.dot_op<{opIdx = 1, parent = #mma1}>>
        %124 = tt.dot %123, %80, %122 {allowTF32 = true} : tensor<128x64xf16, #triton_gpu.dot_op<{opIdx = 0, parent = #mma1}>> * tensor<64x128xf16, #triton_gpu.dot_op<{opIdx = 1, parent = #mma1}>> -> tensor<128x128xf32, #mma1>
        %125 = triton_gpu.convert_layout %124 : (tensor<128x128xf32, #mma1>) -> tensor<128x128xf32, #blocked3>
        %126 = arith.mulf %109, %125 : tensor<128x128xf32, #blocked3>
        %127 = arith.mulf %126, %39 : tensor<128x128xf32, #blocked3>
        %128 = arith.truncf %127 : tensor<128x128xf32, #blocked3> to tensor<128x128xf16, #blocked3>
        %129 = triton_gpu.convert_layout %128 : (tensor<128x128xf16, #blocked3>) -> tensor<128x128xf16, #shared1>
        %130 = tt.trans %129 : (tensor<128x128xf16, #shared1>) -> tensor<128x128xf16, #shared0>
        %131 = triton_gpu.convert_layout %130 : (tensor<128x128xf16, #shared0>) -> tensor<128x128xf16, #triton_gpu.dot_op<{opIdx = 0, parent = #mma0}>>
        %132 = triton_gpu.convert_layout %93 : (tensor<128x64xf16, #blocked1>) -> tensor<128x64xf16, #triton_gpu.dot_op<{opIdx = 1, parent = #mma0}>>
        %133 = tt.dot %131, %132, %arg28 {allowTF32 = true} : tensor<128x128xf16, #triton_gpu.dot_op<{opIdx = 0, parent = #mma0}>> * tensor<128x64xf16, #triton_gpu.dot_op<{opIdx = 1, parent = #mma0}>> -> tensor<128x64xf32, #mma0>
        %134 = tt.load %arg29 {cache = 1 : i32, evict = 1 : i32, isVolatile = false} : tensor<128x64xf32, #blocked2>
        %135 = triton_gpu.convert_layout %134 : (tensor<128x64xf32, #blocked2>) -> tensor<128x64xf32, #mma0>
        %136 = triton_gpu.convert_layout %128 : (tensor<128x128xf16, #blocked3>) -> tensor<128x128xf16, #triton_gpu.dot_op<{opIdx = 0, parent = #mma0}>>
        %81 = triton_gpu.convert_layout %59 : (tensor<128x64xf16, #blocked1>) -> tensor<128x64xf16, #triton_gpu.dot_op<{opIdx = 1, parent = #mma0}>>
        %137 = tt.dot %136, %81, %135 {allowTF32 = true} : tensor<128x128xf16, #triton_gpu.dot_op<{opIdx = 0, parent = #mma0}>> * tensor<128x64xf16, #triton_gpu.dot_op<{opIdx = 1, parent = #mma0}>> -> tensor<128x64xf32, #mma0>
        %138 = triton_gpu.convert_layout %137 : (tensor<128x64xf32, #mma0>) -> tensor<128x64xf32, #blocked2>
        tt.store %arg29, %138 : tensor<128x64xf32, #blocked2>
        %139 = tt.addptr %arg29, %43 : tensor<128x64x!tt.ptr<f32>, #blocked2>, tensor<128x64xi32, #blocked2>
        %140 = tt.addptr %arg30, %42 : tensor<128x64x!tt.ptr<f16>, #blocked1>, tensor<128x64xi32, #blocked1>
        %141 = tt.addptr %arg31, %42 : tensor<128x64x!tt.ptr<f16>, #blocked1>, tensor<128x64xi32, #blocked1>
        scf.yield %116, %133, %139, %140, %141 : tensor<128x64xf32, #mma0>, tensor<128x64xf32, #mma0>, tensor<128x64x!tt.ptr<f32>, #blocked2>, tensor<128x64x!tt.ptr<f16>, #blocked1>, tensor<128x64x!tt.ptr<f16>, #blocked1>
      }
      %83 = triton_gpu.convert_layout %82#1 : (tensor<128x64xf32, #mma0>) -> tensor<128x64xf32, #blocked1>
      %84 = triton_gpu.convert_layout %82#0 : (tensor<128x64xf32, #mma0>) -> tensor<128x64xf32, #blocked1>
      %85 = tt.addptr %44, %62 : tensor<128x64x!tt.ptr<f16>, #blocked1>, tensor<128x64xi32, #blocked1>
      %86 = arith.truncf %84 : tensor<128x64xf32, #blocked1> to tensor<128x64xf16, #blocked1>
      tt.store %85, %86 : tensor<128x64xf16, #blocked1>
      %87 = tt.addptr %45, %57 : tensor<128x64x!tt.ptr<f16>, #blocked1>, tensor<128x64xi32, #blocked1>
      %88 = arith.truncf %83 : tensor<128x64xf32, #blocked1> to tensor<128x64xf16, #blocked1>
      tt.store %87, %88 : tensor<128x64xf16, #blocked1>
    }
    return
  }
}