[Bug]: Model fails on flashinfer_attention_prepare_metadata()

System Info

-H100

Who can help?

No response

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction

build_and_run_ad.py --model apple/OpenELM-450M-Instruct --args.model-factory AutoModelForCausalLM '--args.model-kwargs={}' --args.tokenizer meta-llama/Llama-2-7b-hf --args.world-size 1 --args.compile-backend torch-compile --args.attn-backend flashinfer --args.runtime trtllm --args.skip-loading-weights False --args.transforms.detect-sharding.simple-shard-only False --args.max-seq-len 512 --benchmark.enabled False --benchmark.results-path /jet/logs/basic/auto-deploy-model-coverage_ab-flashinfer_b-false_cb-torch-compile_m-apple-openelm-450m-instruct_mf-automodelforcausallm_mk--_msl-512_r-trtllm_sso-false_sw-false_t-meta-llama-llama-2-7b-hf_ws-1/extra.json --benchmark.store-results true

Expected behavior

model should pass

actual behavior

0: File "/opt/tensorrt-llm/tensorrt_llm/_torch/pyexecutor/py_executor.py", line 2025, in _forward_step
0: outputs = forward(scheduled_requests, self.resource_manager,
0: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
0: File "/opt/tensorrt-llm/tensorrt_llm/_torch/pyexecutor/py_executor.py", line 2013, in forward
0: return self.model_engine.forward(
0: ^^^^^^^^^^^^^^^^^^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/utils/_contextlib.py", line 120, in decorate_context
0: return func(*args, **kwargs)
0: ^^^^^^^^^^^^^^^^^^^^^
0: File "/opt/tensorrt-llm/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py", line 325, in forward
0: logits = self._compute_logits()
0: ^^^^^^^^^^^^^^^^^^^^^^
0: File "/opt/tensorrt-llm/tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py", line 300, in _compute_logits
0: logits: torch.Tensor = self.model(**self.cache_seq_interface.named_args)[0]
0: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/eval_frame.py", line 414, in call
0: return super().call(*args, **kwargs)
0: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
0: return self._call_impl(*args, **kwargs)
0: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1786, in _call_impl
0: return forward_call(*args, **kwargs)
0: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/eval_frame.py", line 832, in compile_wrapper
0: return fn(*args, **kwargs)
0: ^^^^^^^^^^^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/fx/graph_module.py", line 837, in call_wrapped
0: return self._wrapped_call(self, *args, **kwargs)
0: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/fx/graph_module.py", line 413, in call
0: raise e
0: File "/usr/local/lib/python3.12/dist-packages/torch/fx/graph_module.py", line 400, in call
0: return super(self.cls, obj).call(*args, **kwargs) # type: ignore[misc]
0: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
0: return self._call_impl(*args, **kwargs)
0: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1881, in _call_impl
0: return inner()
0: ^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1829, in inner
0: result = forward_call(*args, **kwargs)
0: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
0: File "<eval_with_key>.1350", line 208, in forward
0: flashinfer_attention_prepare_metadata = torch.ops.auto_deploy.flashinfer_attention_prepare_metadata(position_ids, seq_len, input_pos, cache_loc, pages_per_seq, slot_idx, 64, 128); position_ids = seq_len = input_pos = cache_loc = pages_per_seq = slot_idx = None
0: File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/convert_frame.py", line 1872, in call
0: result = self._torchdynamo_orig_backend(
0: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/convert_frame.py", line 1621, in call
0: result = self._inner_convert(
0: ^^^^^^^^^^^^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/convert_frame.py", line 688, in call
0: result = _compile(
0: ^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/convert_frame.py", line 1430, in _compile
0: guarded_code, tracer_output = compile_inner(code, one_graph, hooks)
0: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/_utils_internal.py", line 92, in wrapper_function
0: return function(*args, **kwargs)
0: ^^^^^^^^^^^^^^^^^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/convert_frame.py", line 1114, in compile_inner
0: return _compile_inner(code, one_graph, hooks)
0: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/convert_frame.py", line 1148, in _compile_inner
0: dynamo_output = compile_frame(
0: ^^^^^^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/convert_frame.py", line 1029, in compile_frame
0: bytecode, tracer_output = transform_code_object(code, transform)
0: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/bytecode_transformation.py", line 1592, in transform_code_object
0: tracer_output = transformations(instructions, code_options)
0: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/convert_frame.py", line 1001, in transform
0: tracer_output = trace_frame(
0: ^^^^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/convert_frame.py", line 312, in _fn
0: return fn(*args, **kwargs)
0: ^^^^^^^^^^^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/convert_frame.py", line 815, in trace_frame
0: run_tracer()
0: File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/convert_frame.py", line 797, in run_tracer
0: tracer.run()
0: File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/symbolic_convert.py", line 1487, in run
0: while self.step():
0: ^^^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/symbolic_convert.py", line 1348, in step
0: self.dispatch_table[inst.opcode](self, inst)
0: File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/symbolic_convert.py", line 904, in wrapper
0: return inner_fn(self, inst)
0: ^^^^^^^^^^^^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/symbolic_convert.py", line 3411, in CALL
0: self._call(inst)
0: File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/symbolic_convert.py", line 3405, in _call
0: self.call_function(fn, args, kwargs)
0: File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/symbolic_convert.py", line 1266, in call_function
0: self.push(fn.call_function(self, args, kwargs)) # type: ignore[arg-type]
0: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/variables/torch.py", line 1517, in call_function
0: tensor_variable = wrap_fx_proxy(
0: ^^^^^^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/variables/builder.py", line 2644, in wrap_fx_proxy
0: return wrap_fx_proxy_cls(target_cls=TensorVariable, **kwargs)
0: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/variables/builder.py", line 2710, in wrap_fx_proxy_cls
0: return _wrap_fx_proxy(
0: ^^^^^^^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/variables/builder.py", line 2808, in _wrap_fx_proxy
0: example_value = get_fake_value(proxy.node, tx, allow_non_graph_fake=True)
0: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/utils.py", line 3478, in get_fake_value
0: raise TorchRuntimeError(str(e)).with_traceback(e.traceback) from None
0: File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/utils.py", line 3376, in get_fake_value
0: ret_val = wrap_fake_exception(
0: ^^^^^^^^^^^^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/utils.py", line 2864, in wrap_fake_exception
0: return fn()
0: ^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/utils.py", line 3377, in
0: lambda: run_node(tx.output, node, args, kwargs, nnmodule)
0: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/utils.py", line 3587, in run_node
0: raise RuntimeError(make_error_message(e)).with_traceback(
0: File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/utils.py", line 3546, in run_node
0: return node.target(*args, **kwargs) # type: ignore[operator]
0: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/_ops.py", line 841, in call
0: return self._op(*args, **kwargs)
0: ^^^^^^^^^^^^^^^^^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/utils/_stats.py", line 28, in wrapper
0: return fn(*args, **kwargs)
0: ^^^^^^^^^^^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/_subclasses/fake_tensor.py", line 1376, in torch_dispatch
0: return self.dispatch(func, types, args, kwargs)
0: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/_subclasses/fake_tensor.py", line 2092, in dispatch
0: return self._cached_dispatch_impl(func, types, args, kwargs)
0: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/_subclasses/fake_tensor.py", line 1511, in _cached_dispatch_impl
0: output = self._dispatch_impl(func, types, args, kwargs)
0: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/_subclasses/fake_tensor.py", line 2611, in _dispatch_impl
0: return maybe_propagate_real_tensors(fast_impl(self, *args, **kwargs))
0: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/_subclasses/fake_impls.py", line 1244, in fast_binary_impl
0: return slow("error")
0: ^^^^^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/_subclasses/fake_impls.py", line 1165, in slow
0: return slow_ref(*args, **kwargs)
0: ^^^^^^^^^^^^^^^^^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/_prims_common/wrappers.py", line 309, in _fn
0: result = fn(*args, **kwargs)
0: ^^^^^^^^^^^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/_compile.py", line 53, in inner
0: return disable_fn(*args, **kwargs)
0: ^^^^^^^^^^^^^^^^^^^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/eval_frame.py", line 1044, in _fn
0: return fn(*args, **kwargs)
0: ^^^^^^^^^^^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/_prims_common/wrappers.py", line 149, in _fn
0: result = fn(**bound.arguments)
0: ^^^^^^^^^^^^^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/_refs/init.py", line 1107, in _ref
0: output = prim(a, b)
0: ^^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/_refs/init.py", line 1714, in mul
0: return prims.mul(a, b)
0: ^^^^^^^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/_ops.py", line 841, in call
0: return self._op(*args, **kwargs)
0: ^^^^^^^^^^^^^^^^^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/utils/_stats.py", line 28, in wrapper
0: return fn(*args, **kwargs)
0: ^^^^^^^^^^^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/_subclasses/fake_tensor.py", line 1376, in torch_dispatch
0: return self.dispatch(func, types, args, kwargs)
0: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/_subclasses/fake_tensor.py", line 2092, in dispatch
0: return self._cached_dispatch_impl(func, types, args, kwargs)
0: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/_subclasses/fake_tensor.py", line 1511, in _cached_dispatch_impl
0: output = self._dispatch_impl(func, types, args, kwargs)
0: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
0: File "/usr/local/lib/python3.12/dist-packages/torch/_subclasses/fake_tensor.py", line 2657, in _dispatch_impl
0: func.prim_meta_impl(args, **kwargs)
0: File "/usr/local/lib/python3.12/dist-packages/torch/_prims/init.py", line 404, in prim_elementwise_meta
0: utils.check_same_device(*args, allow_cpu_scalar_tensors=True)
0: File "/usr/local/lib/python3.12/dist-packages/torch/_prims_common/init.py", line 867, in check_same_device
0: raise RuntimeError(msg)
0: torch._dynamo.exc.TorchRuntimeError: Dynamo failed to run FX node with fake tensors: call_function aten.mul.Tensor((FakeTensor(..., device='cuda:0', size=(1, 12, s70, 64)), FakeTensor(..., device='meta', size=(1, 1, s70, 64))), **{}): got RuntimeError('Tensor on device meta is not on the expected device cuda:0!')
0:
0: from user code:
0: File "<eval_with_key>.1350", line 253, in torch_dynamo_resume_in_forward_at_208
0: mul_7 = torch.ops.aten.mul.Tensor(to_5, slice_2); slice_2 = None
0:

additional notes

NA

Before submitting a new issue...

• Make sure you already searched for relevant issues, and checked the documentation and examples for answers to frequently asked questions.