Compiler crash in LLVMCPUSelectLoweringStrategy with dynamic-shape iree_linalg_ext.attention

[stellaraccident]

Description

The compiler crashes during the LLVMCPUSelectLoweringStrategy pass when compiling an iree_linalg_ext.attention op with fully dynamic 3D input shapes (tensor<?x?x?xf32>).

The crash occurs in setRootConfig for AttentionOp when calling limitVectorTileSizes.

Reproducer

Save as attention_crash.mlir:

module @attention_crash {
  util.func public @attention(
      %query: tensor<?x?x?xf32>,
      %key: tensor<?x?x?xf32>,
      %value: tensor<?x?x?xf32>,
      %scale: f32
  ) -> tensor<?x?x?xf32> {
    %c0 = arith.constant 0 : index
    %c1 = arith.constant 1 : index
    %c2 = arith.constant 2 : index

    %batch = tensor.dim %query, %c0 : tensor<?x?x?xf32>
    %seq_len = tensor.dim %query, %c1 : tensor<?x?x?xf32>
    %head_dim = tensor.dim %query, %c2 : tensor<?x?x?xf32>

    %output_init = tensor.empty(%batch, %seq_len, %head_dim) : tensor<?x?x?xf32>
    %output = iree_linalg_ext.attention {
      indexing_maps = [
        affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2)>,
        affine_map<(d0, d1, d2, d3, d4) -> (d0, d3, d2)>,
        affine_map<(d0, d1, d2, d3, d4) -> (d0, d3, d4)>,
        affine_map<(d0, d1, d2, d3, d4) -> ()>,
        affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d4)>
      ]
    } ins(%query, %key, %value, %scale : tensor<?x?x?xf32>, tensor<?x?x?xf32>, tensor<?x?x?xf32>, f32)
      outs(%output_init : tensor<?x?x?xf32>) {
    ^bb0(%arg0: f32):
      iree_linalg_ext.yield %arg0 : f32
    } -> tensor<?x?x?xf32>

    util.return %output : tensor<?x?x?xf32>
  }
}

Compile with:

iree-compile attention_crash.mlir --iree-hal-target-backends=llvm-cpu -o /dev/null

Stack trace (key frames)

#4 limitVectorTileSizes(...) KernelDispatch.cpp:747
#8 setRootConfig(..., AttentionOp) KernelDispatch.cpp:2344
#14 LLVMCPUSelectLoweringStrategyPass::runOnOperation() LLVMCPUSelectLoweringStrategy.cpp:238

Environment

• IREE: pip package (current release)
• OS: Fedora 43, Linux 6.17
• Target: llvm-cpu

Notes

Static shapes compile successfully. The crash appears to occur during tile size computation when shapes are unknown at compile time.

---

This issue was authored by Claude (Opus 4.5) and reviewed by @stellaraccident.

[stellaraccident]

The SIGFPE crash is fixed, but dynamic shapes + attention is still broken. Compilation now "succeeds" but produces garbage:

$ iree-compile attention_crash.mlir --iree-hal-target-backends=llvm-cpu -o /tmp/out.vmfb
warning: -:1:0: stack frame size (144115188075857464) exceeds limit (4294967295) in function 'attention_dispatch_3_attention_DxDxDxDxD'

That's a 144 petabyte stack frame - the dynamic dimensions are being treated as garbage values somewhere in the lowering.

Running the compiled module segfaults.

---

Comment authored by Claude (Opus 4.5) and reviewed by @stellaraccident.

[stellaraccident]

@hanhanW still needs work. PTAL

[stellaraccident]

Additional failure mode: memref.alloca hoisting error

When iree_linalg_ext.attention with dynamic shapes is called from within another function (via util.call), the error manifests differently as a memref.alloca hoisting failure rather than the 144PB stack frame warning.

Reproducer (linked attention block)

This IR was produced by iree-link composing an attention block that calls the attention op:

#map = affine_map<(d0) -> (d0)>
#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, 0)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, 1)>
#map3 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1)>
#map4 = affine_map<(d0, d1, d2, d3, d4) -> (d3)>
#map5 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>
#map6 = affine_map<(d0, d1, d2, d3) -> (d0, d2, d1, d3)>
#map7 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
#map8 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2)>
#map9 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d3, d2)>
#map10 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d3, d4)>
#map11 = affine_map<(d0, d1, d2, d3, d4) -> ()>
#map12 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d4)>
#map13 = affine_map<(d0, d1, d2) -> (d1, d2)>
#map14 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
#map15 = affine_map<(d0, d1, d2) -> (d2)>
module @attention_block_components {
  util.func private @position_components.rope(%arg0: tensor<?x?x?x?xf32>, %arg1: tensor<?x?xi64>, %arg2: f32, %arg3: f32) -> tensor<?x?x?x?xf32> {
    %c0 = arith.constant 0 : index
    %c1 = arith.constant 1 : index
    %c2 = arith.constant 2 : index
    %c3 = arith.constant 3 : index
    %dim = tensor.dim %arg0, %c0 : tensor<?x?x?x?xf32>
    %dim_0 = tensor.dim %arg0, %c1 : tensor<?x?x?x?xf32>
    %dim_1 = tensor.dim %arg0, %c2 : tensor<?x?x?x?xf32>
    %dim_2 = tensor.dim %arg0, %c3 : tensor<?x?x?x?xf32>
    %0 = arith.divsi %dim_2, %c2 : index
    %expanded = tensor.expand_shape %arg0 [[0], [1], [2], [3, 4]] output_shape [%dim, %dim_0, %dim_1, %0, 2] : tensor<?x?x?x?xf32> into tensor<?x?x?x?x2xf32>
    %1 = arith.index_cast %dim_2 : index to i32
    %2 = arith.sitofp %1 : i32 to f32
    %3 = tensor.empty(%0) : tensor<?xf32>
    %4 = linalg.generic {indexing_maps = [#map], iterator_types = ["parallel"]} outs(%3 : tensor<?xf32>) {
    ^bb0(%out: f32):
      %7 = linalg.index 0 : index
      %8 = arith.index_cast %7 : index to i32
      %9 = arith.sitofp %8 : i32 to f32
      %cst = arith.constant 2.000000e+00 : f32
      %10 = arith.mulf %cst, %9 : f32
      %11 = arith.divf %10, %2 : f32
      %12 = arith.negf %11 : f32
      %13 = math.powf %arg2, %12 : f32
      %14 = arith.mulf %13, %arg3 : f32
      linalg.yield %14 : f32
    } -> tensor<?xf32>
    %5 = tensor.empty(%dim, %dim_0, %dim_1, %0) : tensor<?x?x?x?x2xf32>
    %6 = linalg.generic {indexing_maps = [#map1, #map2, #map3, #map4, #map5], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%expanded, %expanded, %arg1, %4 : tensor<?x?x?x?x2xf32>, tensor<?x?x?x?x2xf32>, tensor<?x?xi64>, tensor<?xf32>) outs(%5 : tensor<?x?x?x?x2xf32>) {
    ^bb0(%in: f32, %in_3: f32, %in_4: i64, %in_5: f32, %out: f32):
      %7 = arith.trunci %in_4 : i64 to i32
      %8 = arith.sitofp %7 : i32 to f32
      %9 = arith.mulf %8, %in_5 : f32
      %10 = math.cos %9 : f32
      %11 = math.sin %9 : f32
      %12 = linalg.index 4 : index
      %c0_6 = arith.constant 0 : index
      %13 = arith.cmpi eq, %12, %c0_6 : index
      %14 = scf.if %13 -> (f32) {
        %15 = arith.mulf %in, %10 : f32
        %16 = arith.mulf %in_3, %11 : f32
        %17 = arith.subf %15, %16 : f32
        scf.yield %17 : f32
      } else {
        %15 = arith.mulf %in, %11 : f32
        %16 = arith.mulf %in_3, %10 : f32
        %17 = arith.addf %15, %16 : f32
        scf.yield %17 : f32
      }
      linalg.yield %14 : f32
    } -> tensor<?x?x?x?x2xf32>
    %collapsed = tensor.collapse_shape %6 [[0], [1], [2], [3, 4]] : tensor<?x?x?x?x2xf32> into tensor<?x?x?x?xf32>
    util.return %collapsed : tensor<?x?x?x?xf32>
  }
  util.func private @attention_components.attention_gqa(%arg0: tensor<?x?x?x?xf32>, %arg1: tensor<?x?x?x?xf32>, %arg2: tensor<?x?x?x?xf32>, %arg3: f32) -> tensor<?x?x?x?xf32> {
    %c0 = arith.constant 0 : index
    %c1 = arith.constant 1 : index
    %c2 = arith.constant 2 : index
    %c3 = arith.constant 3 : index
    %dim = tensor.dim %arg0, %c0 : tensor<?x?x?x?xf32>
    %dim_0 = tensor.dim %arg0, %c1 : tensor<?x?x?x?xf32>
    %dim_1 = tensor.dim %arg0, %c2 : tensor<?x?x?x?xf32>
    %dim_2 = tensor.dim %arg0, %c3 : tensor<?x?x?x?xf32>
    %0 = tensor.empty(%dim, %dim_1, %dim_0, %dim_2) : tensor<?x?x?x?xf32>
    %1 = linalg.generic {indexing_maps = [#map6, #map7], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg0 : tensor<?x?x?x?xf32>) outs(%0 : tensor<?x?x?x?xf32>) {
    ^bb0(%in: f32, %out: f32):
      linalg.yield %in : f32
    } -> tensor<?x?x?x?xf32>
    %2 = tensor.empty(%dim, %dim_1, %dim_0, %dim_2) : tensor<?x?x?x?xf32>
    %3 = linalg.generic {indexing_maps = [#map6, #map7], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg1 : tensor<?x?x?x?xf32>) outs(%2 : tensor<?x?x?x?xf32>) {
    ^bb0(%in: f32, %out: f32):
      linalg.yield %in : f32
    } -> tensor<?x?x?x?xf32>
    %4 = tensor.empty(%dim, %dim_1, %dim_0, %dim_2) : tensor<?x?x?x?xf32>
    %5 = linalg.generic {indexing_maps = [#map6, #map7], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg2 : tensor<?x?x?x?xf32>) outs(%4 : tensor<?x?x?x?xf32>) {
    ^bb0(%in: f32, %out: f32):
      linalg.yield %in : f32
    } -> tensor<?x?x?x?xf32>
    %6 = arith.muli %dim, %dim_1 : index
    %collapsed = tensor.collapse_shape %1 [[0, 1], [2], [3]] : tensor<?x?x?x?xf32> into tensor<?x?x?xf32>
    %collapsed_3 = tensor.collapse_shape %3 [[0, 1], [2], [3]] : tensor<?x?x?x?xf32> into tensor<?x?x?xf32>
    %collapsed_4 = tensor.collapse_shape %5 [[0, 1], [2], [3]] : tensor<?x?x?x?xf32> into tensor<?x?x?xf32>
    %7 = tensor.empty(%6, %dim_0, %dim_2) : tensor<?x?x?xf32>
    %8 = iree_linalg_ext.attention {indexing_maps = [#map8, #map9, #map10, #map11, #map12]} ins(%collapsed, %collapsed_3, %collapsed_4, %arg3 : tensor<?x?x?xf32>, tensor<?x?x?xf32>, tensor<?x?x?xf32>, f32) outs(%7 : tensor<?x?x?xf32>) {
    ^bb0(%arg4: f32):
      iree_linalg_ext.yield %arg4 : f32
    } -> tensor<?x?x?xf32>
    %expanded = tensor.expand_shape %8 [[0, 1], [2], [3]] output_shape [%dim, %dim_1, %dim_0, %dim_2] : tensor<?x?x?xf32> into tensor<?x?x?x?xf32>
    %9 = tensor.empty(%dim, %dim_0, %dim_1, %dim_2) : tensor<?x?x?x?xf32>
    %10 = linalg.generic {indexing_maps = [#map6, #map7], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded : tensor<?x?x?x?xf32>) outs(%9 : tensor<?x?x?x?xf32>) {
    ^bb0(%in: f32, %out: f32):
      linalg.yield %in : f32
    } -> tensor<?x?x?x?xf32>
    util.return %10 : tensor<?x?x?x?xf32>
  }
  util.func public @attention_block(%arg0: tensor<?x?x?xf32>, %arg1: tensor<?x?xi64>, %arg2: tensor<?x?xf32>, %arg3: tensor<?x?xf32>, %arg4: tensor<?x?xf32>, %arg5: tensor<?x?xf32>, %arg6: tensor<?xf32>, %arg7: tensor<?xf32>, %arg8: tensor<?xf32>, %arg9: tensor<?xf32>, %arg10: i1, %arg11: index, %arg12: index, %arg13: index, %arg14: f32, %arg15: f32) -> tensor<?x?x?xf32> {
    %c0 = arith.constant 0 : index
    %c1 = arith.constant 1 : index
    %dim = tensor.dim %arg0, %c0 : tensor<?x?x?xf32>
    %dim_0 = tensor.dim %arg0, %c1 : tensor<?x?x?xf32>
    %dim_1 = tensor.dim %arg3, %c1 : tensor<?x?xf32>
    %dim_2 = tensor.dim %arg2, %c0 : tensor<?x?xf32>
    %dim_3 = tensor.dim %arg3, %c0 : tensor<?x?xf32>
    %dim_4 = tensor.dim %arg4, %c0 : tensor<?x?xf32>
    %dim_5 = tensor.dim %arg5, %c0 : tensor<?x?xf32>
    %0 = tensor.empty(%dim, %dim_2, %arg13) : tensor<?x?x?xf32>
    %1 = linalg.generic {indexing_maps = [#map13, #map14], iterator_types = ["parallel", "parallel", "parallel"]} ins(%arg2 : tensor<?x?xf32>) outs(%0 : tensor<?x?x?xf32>) {
    ^bb0(%in: f32, %out: f32):
      linalg.yield %in : f32
    } -> tensor<?x?x?xf32>
    %2 = tensor.empty(%dim, %dim_3, %dim_1) : tensor<?x?x?xf32>
    %3 = linalg.generic {indexing_maps = [#map13, #map14], iterator_types = ["parallel", "parallel", "parallel"]} ins(%arg3 : tensor<?x?xf32>) outs(%2 : tensor<?x?x?xf32>) {
    ^bb0(%in: f32, %out: f32):
      linalg.yield %in : f32
    } -> tensor<?x?x?xf32>
    %4 = tensor.empty(%dim, %dim_4, %dim_1) : tensor<?x?x?xf32>
    %5 = linalg.generic {indexing_maps = [#map13, #map14], iterator_types = ["parallel", "parallel", "parallel"]} ins(%arg4 : tensor<?x?xf32>) outs(%4 : tensor<?x?x?xf32>) {
    ^bb0(%in: f32, %out: f32):
      linalg.yield %in : f32
    } -> tensor<?x?x?xf32>
    %cst = arith.constant 0.000000e+00 : f32
    %6 = tensor.empty(%dim, %dim_0, %arg13) : tensor<?x?x?xf32>
    %7 = linalg.fill ins(%cst : f32) outs(%6 : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
    %8 = linalg.batch_matmul ins(%arg0, %1 : tensor<?x?x?xf32>, tensor<?x?x?xf32>) outs(%7 : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
    %9 = tensor.empty(%dim, %dim_0, %dim_1) : tensor<?x?x?xf32>
    %10 = linalg.fill ins(%cst : f32) outs(%9 : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
    %11 = linalg.batch_matmul ins(%arg0, %3 : tensor<?x?x?xf32>, tensor<?x?x?xf32>) outs(%10 : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
    %12 = tensor.empty(%dim, %dim_0, %dim_1) : tensor<?x?x?xf32>
    %13 = linalg.fill ins(%cst : f32) outs(%12 : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
    %14 = linalg.batch_matmul ins(%arg0, %5 : tensor<?x?x?xf32>, tensor<?x?x?xf32>) outs(%13 : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
    %15 = scf.if %arg10 -> (tensor<?x?x?xf32>) {
      %32 = linalg.generic {indexing_maps = [#map14, #map15, #map14], iterator_types = ["parallel", "parallel", "parallel"]} ins(%8, %arg6 : tensor<?x?x?xf32>, tensor<?xf32>) outs(%6 : tensor<?x?x?xf32>) {
      ^bb0(%in: f32, %in_8: f32, %out: f32):
        %33 = arith.addf %in, %in_8 : f32
        linalg.yield %33 : f32
      } -> tensor<?x?x?xf32>
      scf.yield %32 : tensor<?x?x?xf32>
    } else {
      scf.yield %8 : tensor<?x?x?xf32>
    }
    %16 = scf.if %arg10 -> (tensor<?x?x?xf32>) {
      %32 = linalg.generic {indexing_maps = [#map14, #map15, #map14], iterator_types = ["parallel", "parallel", "parallel"]} ins(%11, %arg7 : tensor<?x?x?xf32>, tensor<?xf32>) outs(%9 : tensor<?x?x?xf32>) {
      ^bb0(%in: f32, %in_8: f32, %out: f32):
        %33 = arith.addf %in, %in_8 : f32
        linalg.yield %33 : f32
      } -> tensor<?x?x?xf32>
      scf.yield %32 : tensor<?x?x?xf32>
    } else {
      scf.yield %11 : tensor<?x?x?xf32>
    }
    %17 = scf.if %arg10 -> (tensor<?x?x?xf32>) {
      %32 = linalg.generic {indexing_maps = [#map14, #map15, #map14], iterator_types = ["parallel", "parallel", "parallel"]} ins(%14, %arg8 : tensor<?x?x?xf32>, tensor<?xf32>) outs(%12 : tensor<?x?x?xf32>) {
      ^bb0(%in: f32, %in_8: f32, %out: f32):
        %33 = arith.addf %in, %in_8 : f32
        linalg.yield %33 : f32
      } -> tensor<?x?x?xf32>
      scf.yield %32 : tensor<?x?x?xf32>
    } else {
      scf.yield %14 : tensor<?x?x?xf32>
    }
    %18 = arith.divsi %arg13, %arg11 : index
    %19 = arith.divsi %dim_1, %arg12 : index
    %expanded = tensor.expand_shape %15 [[0], [1], [2, 3]] output_shape [%dim, %dim_0, %arg11, %18] : tensor<?x?x?xf32> into tensor<?x?x?x?xf32>
    %expanded_6 = tensor.expand_shape %16 [[0], [1], [2, 3]] output_shape [%dim, %dim_0, %arg12, %19] : tensor<?x?x?xf32> into tensor<?x?x?x?xf32>
    %expanded_7 = tensor.expand_shape %17 [[0], [1], [2, 3]] output_shape [%dim, %dim_0, %arg12, %19] : tensor<?x?x?xf32> into tensor<?x?x?x?xf32>
    %20 = util.call @position_components.rope(%expanded, %arg1, %arg14, %arg15) : (tensor<?x?x?x?xf32>, tensor<?x?xi64>, f32, f32) -> tensor<?x?x?x?xf32>
    %21 = util.call @position_components.rope(%expanded_6, %arg1, %arg14, %arg15) : (tensor<?x?x?x?xf32>, tensor<?x?xi64>, f32, f32) -> tensor<?x?x?x?xf32>
    %22 = arith.index_cast %18 : index to i32
    %23 = arith.sitofp %22 : i32 to f32
    %24 = math.rsqrt %23 : f32
    %25 = util.call @attention_components.attention_gqa(%20, %21, %expanded_7, %24) : (tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>, f32) -> tensor<?x?x?x?xf32>
    %collapsed = tensor.collapse_shape %25 [[0], [1], [2, 3]] : tensor<?x?x?x?xf32> into tensor<?x?x?xf32>
    %26 = tensor.empty(%dim, %dim_5, %arg13) : tensor<?x?x?xf32>
    %27 = linalg.generic {indexing_maps = [#map13, #map14], iterator_types = ["parallel", "parallel", "parallel"]} ins(%arg5 : tensor<?x?xf32>) outs(%26 : tensor<?x?x?xf32>) {
    ^bb0(%in: f32, %out: f32):
      linalg.yield %in : f32
    } -> tensor<?x?x?xf32>
    %28 = tensor.empty(%dim, %dim_0, %arg13) : tensor<?x?x?xf32>
    %29 = linalg.fill ins(%cst : f32) outs(%28 : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
    %30 = linalg.batch_matmul ins(%collapsed, %27 : tensor<?x?x?xf32>, tensor<?x?x?xf32>) outs(%29 : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
    %31 = scf.if %arg10 -> (tensor<?x?x?xf32>) {
      %32 = linalg.generic {indexing_maps = [#map14, #map15, #map14], iterator_types = ["parallel", "parallel", "parallel"]} ins(%30, %arg9 : tensor<?x?x?xf32>, tensor<?xf32>) outs(%28 : tensor<?x?x?xf32>) {
      ^bb0(%in: f32, %in_8: f32, %out: f32):
        %33 = arith.addf %in, %in_8 : f32
        linalg.yield %33 : f32
      } -> tensor<?x?x?xf32>
      scf.yield %32 : tensor<?x?x?xf32>
    } else {
      scf.yield %30 : tensor<?x?x?xf32>
    }
    util.return %31 : tensor<?x?x?xf32>
  }
}

Compile command

iree-compile attention_block_linked.mlir --iree-hal-target-backends=llvm-cpu -o out.vmfb

Error output

attention_block_linked.mlir:101:10: error: 'memref.alloca' op all stack allocations need to be hoisted to the entry block of the function
    %8 = iree_linalg_ext.attention {indexing_maps = [#map8, #map9, #map10, #map11, #map12]} ins(%collapsed, %collapsed_3, %collapsed_4, %arg3 : tensor<?x?x?xf32>, tensor<?x?x?xf32>, tensor<?x?x?xf32>, f32) outs(%7 : tensor<?x?x?xf32>) {
         ^
attention_block_linked.mlir:188:11: note: called from
    %25 = util.call @attention_components.attention_gqa(%20, %21, %expanded_7, %24) : (tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>, f32) -> tensor<?x?x?x?xf32>
          ^
attention_block_linked.mlir:113:3: note: called from
  util.func public @attention_block(...) -> tensor<?x?x?xf32> {
  ^
attention_block_linked.mlir:101:10: note: see current operation: %171 = "memref.alloca"(%166, %104) <{alignment = 64 : i64, operandSegmentSizes = array<i32: 2, 0>}> : (index, index) -> memref<1x?x?xf32>

Summary

Two failure modes observed with dynamic-shape attention on CPU:

1. Standalone attention op: 144PB stack frame warning, compiles but segfaults at runtime
2. Attention called via util.call: memref.alloca hoisting error, fails to compile

Both stem from the same underlying issue with dynamic shapes in the attention op lowering.

[stellaraccident]

Workaround: partial dynamism with static head dims

For anyone hitting this — making n_head and head_dim static while keeping batch and seq_len dynamic compiles and runs correctly. This narrows the issue to fully dynamic shapes specifically.

Working wrapper (partial dynamism)

The wrapper has tensor<?x?x8x16xf32> (dynamic batch/seq, static head dims) and casts to fully dynamic for the callee. util.assume.int constrains seq_len divisibility:

module @test_wrapper {
  util.func private @attention_components.attention_gqa(
      tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>, f32
  ) -> tensor<?x?x?x?xf32>

  util.func public @main(
      %q: tensor<?x?x8x16xf32>, %k: tensor<?x?x8x16xf32>,
      %v: tensor<?x?x8x16xf32>, %scale: f32
  ) -> tensor<?x?x8x16xf32> {
    %c0 = arith.constant 0 : index
    %c1 = arith.constant 1 : index
    %batch_raw = tensor.dim %q, %c0 : tensor<?x?x8x16xf32>
    %seq_raw = tensor.dim %q, %c1 : tensor<?x?x8x16xf32>
    %batch = util.assume.int %batch_raw<umin=1, umax=64> : index
    %seq = util.assume.int %seq_raw<umin=32, umax=4096, udiv=32> : index

    %q_dyn = tensor.cast %q : tensor<?x?x8x16xf32> to tensor<?x?x?x?xf32>
    %k_dyn = tensor.cast %k : tensor<?x?x8x16xf32> to tensor<?x?x?x?xf32>
    %v_dyn = tensor.cast %v : tensor<?x?x8x16xf32> to tensor<?x?x?x?xf32>
    %out_dyn = util.call @attention_components.attention_gqa(
        %q_dyn, %k_dyn, %v_dyn, %scale)
        : (tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>, f32)
        -> tensor<?x?x?x?xf32>
    %out = tensor.cast %out_dyn : tensor<?x?x?x?xf32> to tensor<?x?x8x16xf32>
    util.return %out : tensor<?x?x8x16xf32>
  }
}

Post iree-link IR (what the compiler actually sees)

After linking with attention_gqa.mlir, the callee is inlined. The compiler sees util.assume.int constraints and static dims from the tensor.cast chain:

#map = affine_map<(d0, d1, d2, d3) -> (d0, d2, d1, d3)>
#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
#map2 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2)>
#map3 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d3, d2)>
#map4 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d3, d4)>
#map5 = affine_map<(d0, d1, d2, d3, d4) -> ()>
#map6 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d4)>
module @test_wrapper {
  util.func private @attention_components.attention_gqa(%arg0: tensor<?x?x?x?xf32>, %arg1: tensor<?x?x?x?xf32>, %arg2: tensor<?x?x?x?xf32>, %arg3: f32) -> tensor<?x?x?x?xf32> {
    %c0 = arith.constant 0 : index
    %c1 = arith.constant 1 : index
    %c2 = arith.constant 2 : index
    %c3 = arith.constant 3 : index
    %dim = tensor.dim %arg0, %c0 : tensor<?x?x?x?xf32>
    %dim_0 = tensor.dim %arg0, %c1 : tensor<?x?x?x?xf32>
    %dim_1 = tensor.dim %arg0, %c2 : tensor<?x?x?x?xf32>
    %dim_2 = tensor.dim %arg0, %c3 : tensor<?x?x?x?xf32>
    %0 = tensor.empty(%dim, %dim_1, %dim_0, %dim_2) : tensor<?x?x?x?xf32>
    %1 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg0 : tensor<?x?x?x?xf32>) outs(%0 : tensor<?x?x?x?xf32>) {
    ^bb0(%in: f32, %out: f32):
      linalg.yield %in : f32
    } -> tensor<?x?x?x?xf32>
    %2 = tensor.empty(%dim, %dim_1, %dim_0, %dim_2) : tensor<?x?x?x?xf32>
    %3 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg1 : tensor<?x?x?x?xf32>) outs(%2 : tensor<?x?x?x?xf32>) {
    ^bb0(%in: f32, %out: f32):
      linalg.yield %in : f32
    } -> tensor<?x?x?x?xf32>
    %4 = tensor.empty(%dim, %dim_1, %dim_0, %dim_2) : tensor<?x?x?x?xf32>
    %5 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg2 : tensor<?x?x?x?xf32>) outs(%4 : tensor<?x?x?x?xf32>) {
    ^bb0(%in: f32, %out: f32):
      linalg.yield %in : f32
    } -> tensor<?x?x?x?xf32>
    %6 = arith.muli %dim, %dim_1 : index
    %collapsed = tensor.collapse_shape %1 [[0, 1], [2], [3]] : tensor<?x?x?x?xf32> into tensor<?x?x?xf32>
    %collapsed_3 = tensor.collapse_shape %3 [[0, 1], [2], [3]] : tensor<?x?x?x?xf32> into tensor<?x?x?xf32>
    %collapsed_4 = tensor.collapse_shape %5 [[0, 1], [2], [3]] : tensor<?x?x?x?xf32> into tensor<?x?x?xf32>
    %7 = tensor.empty(%6, %dim_0, %dim_2) : tensor<?x?x?xf32>
    %8 = iree_linalg_ext.attention {indexing_maps = [#map2, #map3, #map4, #map5, #map6]} ins(%collapsed, %collapsed_3, %collapsed_4, %arg3 : tensor<?x?x?xf32>, tensor<?x?x?xf32>, tensor<?x?x?xf32>, f32) outs(%7 : tensor<?x?x?xf32>) {
    ^bb0(%arg4: f32):
      iree_linalg_ext.yield %arg4 : f32
    } -> tensor<?x?x?xf32>
    %expanded = tensor.expand_shape %8 [[0, 1], [2], [3]] output_shape [%dim, %dim_1, %dim_0, %dim_2] : tensor<?x?x?xf32> into tensor<?x?x?x?xf32>
    %9 = tensor.empty(%dim, %dim_0, %dim_1, %dim_2) : tensor<?x?x?x?xf32>
    %10 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded : tensor<?x?x?x?xf32>) outs(%9 : tensor<?x?x?x?xf32>) {
    ^bb0(%in: f32, %out: f32):
      linalg.yield %in : f32
    } -> tensor<?x?x?x?xf32>
    util.return %10 : tensor<?x?x?x?xf32>
  }
  util.func public @main(%arg0: tensor<?x?x8x16xf32>, %arg1: tensor<?x?x8x16xf32>, %arg2: tensor<?x?x8x16xf32>, %arg3: f32) -> tensor<?x?x8x16xf32> {
    %c0 = arith.constant 0 : index
    %c1 = arith.constant 1 : index
    %dim = tensor.dim %arg0, %c0 : tensor<?x?x8x16xf32>
    %dim_0 = tensor.dim %arg0, %c1 : tensor<?x?x8x16xf32>
    %0 = util.assume.int %dim<umin = 1, umax = 64> : index
    %1 = util.assume.int %dim_0<umin = 32, umax = 4096, udiv = 32> : index
    %cast = tensor.cast %arg0 : tensor<?x?x8x16xf32> to tensor<?x?x?x?xf32>
    %cast_1 = tensor.cast %arg1 : tensor<?x?x8x16xf32> to tensor<?x?x?x?xf32>
    %cast_2 = tensor.cast %arg2 : tensor<?x?x8x16xf32> to tensor<?x?x?x?xf32>
    %2 = util.call @attention_components.attention_gqa(%cast, %cast_1, %cast_2, %arg3) : (tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>, f32) -> tensor<?x?x?x?xf32>
    %cast_3 = tensor.cast %2 : tensor<?x?x?x?xf32> to tensor<?x?x8x16xf32>
    util.return %cast_3 : tensor<?x?x8x16xf32>
  }
}

This compiles cleanly with --iree-hal-target-backends=llvm-cpu and produces correct results for multiple batch/seq_len combinations at runtime (tested with batch ∈ {1,2,3,4}, seq_len ∈ {32,64,128}).

The key difference from the fully-dynamic case: n_head and head_dim are known statically through the tensor.cast chain (tensor<?x?x8x16xf32> → tensor<?x?x?x?xf32>), which gives the compiler enough information to size the alloca statically.