CustomCallOp unsupported target: TopK on Llama 3.1 with topk=2

[steeve]

Hi,

Our (@zml) Llama implementation fails to compile if we run with topk > 1.
We're not sure what triggers it, it looks like some pattern matching.

Please find attached our implementation with topk=2.

Thanks!!

error(pjrt): [PJRT_Client_Compile] RunNeuronCCImpl: error condition error != 0: <class 'subprocess.CalledProcessError'>: Command '['neuronx-cc', 'compile', '--framework=XLA', '--target=trn1', '--verbose=35', '--output=/tmp/tmpqle33ids/file.neff', '/tmp/tmpqle33ids/file.code', '--model-type=transformer', '--auto-cast=none']' returned non-zero exit status 70.

llama.mlir.txt