failed to translate executables for sharkai models

[ziereis]

What happened?

cant compile 3b open llama model.

cuda:

failed to translate executables
failed to translate executables
shark_models/shark_llama_3b_bs1.mlir:1121:12: error: failed to run translation of source executable to target executable for backend #hal.executable.target<"cuda", "cuda-nvptx-fb", {iree.gpu.target = #iree_gpu.target<arch = "sm_80", features = "+ptx76", wgp = <compute =  fp64|fp32|fp16|int64|int32|int16|int8, storage =  b64|b32|b16|b8, subgroup =  shuffle|arithmetic, dot =  dp4xi8toi32, mma = [<NV_WMMA_F32_16x16x16_F16>, <NV_WMMA_F16_16x16x16_F16>], subgroup_size_choices = [32], max_workgroup_sizes = [1024, 1024, 1024], max_thread_count_per_workgroup = 1024, max_workgroup_memory_bytes = 166912, max_workgroup_counts = [2147483647, 65535, 65535]>>}>
    %378 = torch.aten.transpose.int %377#0, %int1_170, %int2_171 : !torch.vtensor<[1,32,?,100],f16>, !torch.int, !torch.int -> !torch.vtensor<[1,?,32,100],f16>

for hip the compilation seems to be stuck somewhere, didn't finish compilation after roughly 20 minutes.

Steps to reproduce your issue

model exported with:

python -m sharktank.examples.export_paged_llm_v1 --hf-dataset=SlyEcho/open_llama_3b_v2_f16_gguf --bs=1

failed compilation:

iree-compile --iree-hal-target-device=cuda -iree-cuda-target=sm_80 

compilation either takes very long or is stuck somewhere, tried running it for 20 minutes:

iree-compile --iree-hal-target-device=hip --iree-hip-target=gfx90a

What component(s) does this issue relate to?

Compiler

Version information

4693b1c

Additional context

No response