[Bug]: Model fails on triton_attention_prepare_fused_mha_metadata

System Info

• H100

Who can help?

Error log:

0: Traceback (most recent call last):
0:   File "/usr/lib/python3.12/multiprocessing/process.py", line 314, in _bootstrap
0:     self.run()
0:   File "/usr/lib/python3.12/multiprocessing/process.py", line 108, in run
0:     self._target(*self._args, **self._kwargs)
0:   File "/opt/tensorrt-llm/tensorrt_llm/_torch/auto_deploy/distributed/common.py", line 173, in init_and_run_process
0:     raise e
0:   File "/opt/tensorrt-llm/tensorrt_llm/_torch/auto_deploy/distributed/common.py", line 166, in init_and_run_process
0:     job(rank, size, **kwargs)
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/demollm.py", line 318, in _run_engine
0:     outs = engine.generate_tokens_batched(request_list)
0:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
0:   File "/opt/tensorrt-llm/tensorrt_llm/_torch/auto_deploy/shim/demollm.py", line 181, in generate_tokens_batched
0:     _generate_single_step(i)
0:   File "/opt/tensorrt-llm/tensorrt_llm/_torch/auto_deploy/shim/demollm.py", line 135, in _generate_single_step
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>.4733", line 289, in forward
0:     triton_attention_prepare_fused_mha_metadata = torch.ops.auto_deploy.triton_attention_prepare_fused_mha_metadata(position_ids, seq_len, input_pos, cache_loc, pages_per_seq, slot_idx, 2176, 128);  seq_len = input_pos = cache_loc = pages_per_seq = slot_idx = None
0:   File "<eval_with_key>.4733", line 289, in torch_dynamo_resume_in_forward_at_289
0:     triton_attention_prepare_fused_mha_metadata = torch.ops.auto_deploy.triton_attention_prepare_fused_mha_metadata(position_ids, seq_len, input_pos, cache_loc, pages_per_seq, slot_idx, 2176, 128);  seq_len = input_pos = cache_loc = pages_per_seq = slot_idx = None
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/_functorch/aot_autograd.py", line 1130, in forward
0:     return compiled_fn(full_args)
0:            ^^^^^^^^^^^^^^^^^^^^^^
0:   File "/usr/local/lib/python3.12/dist-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 353, in runtime_wrapper
0:     all_outs = call_func_at_runtime_with_args(
0:                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
0:   File "/usr/local/lib/python3.12/dist-packages/torch/_functorch/_aot_autograd/utils.py", line 129, in call_func_at_runtime_with_args
0:     out = normalize_as_list(f(args))
0:                             ^^^^^^^
0:   File "/usr/local/lib/python3.12/dist-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 724, in inner_fn
0:     outs = compiled_fn(args)
0:            ^^^^^^^^^^^^^^^^^
0:   File "/usr/local/lib/python3.12/dist-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 526, in wrapper
0:     return compiled_fn(runtime_args)
0:            ^^^^^^^^^^^^^^^^^^^^^^^^^
0:   File "/usr/local/lib/python3.12/dist-packages/torch/_inductor/output_code.py", line 611, in __call__
0:     return self.current_callable(inputs)
0:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
0:   File "/usr/local/lib/python3.12/dist-packages/torch/_inductor/utils.py", line 2960, in run
0:     out = model(new_inputs)
0:           ^^^^^^^^^^^^^^^^^
0:   File "/tmp/torchinductor_svc-jet/2s/c2sngro65zhfp7cv5u6nxzpmlqeha7exwe4djxoynxo73t4hajvk.py", line 1143, in call
0:     buf7 = torch.ops.prims.device_put.default(arg6_1, device(type='cuda', index=0))
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/_prims/__init__.py", line 303, in _backend_select_impl
0:     return _prim_impl(*args, **kwargs)
0:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^
0:   File "/usr/local/lib/python3.12/dist-packages/torch/_prims/__init__.py", line 288, in _prim_impl
0:     return impl_aten(*args, **kwargs)
0:            ^^^^^^^^^^^^^^^^^^^^^^^^^^
0:   File "/usr/local/lib/python3.12/dist-packages/torch/_prims/__init__.py", line 1973, in _device_put_aten
0:     return a.to(device, non_blocking=non_blocking)
0:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
0: NotImplementedError: Cannot copy out of meta tensor; no data!

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 microsoft/Phi-3-medium-128k-instruct --args.model-factory AutoModelForCausalLM '--args.model-kwargs={}' --args.tokenizer null --args.world-size 4 --args.compile-backend torch-compile --args.attn-backend triton --args.runtime demollm --args.skip-loading-weights False --args.transforms.detect-sharding.simple-shard-only False --args.max-seq-len 512 --benchmark.enabled True --benchmark.results-path /jet/logs/basic/auto-deploy-model-coverage_ab-triton_b-true_cb-torch-compile_m-microsoft-phi-3-medium-128k-instruct_mf-automodelforcausallm_mk--_msl-512_r-demollm_sso-false_sw-false_t-null_ws-4/extra.json --benchmark.store-results true

Expected behavior

model should pass

actual behavior

model fails on
0: NotImplementedError: Cannot copy out of meta tensor; no data!

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.