[Torchflow] Can't run torchflow on llamaforcausallm with fp32. rmsnorm(at::Tensor&, at::Tensor&, at::Tensor&, double, bool)::<lambda()> failed to dispatch data type Float

System Info

• H100
• weights are in fp32, but I am happy to run it in fp16, bf16 or fp8.

Failing with the same error:

• https://huggingface.co/canopylabs/orpheus-3b-0.1-ft
• https://huggingface.co/baseten/orpheus-3b-0.1-ft
• https://huggingface.co/baseten/orpheus-3b-0.1-ft-fp8/blob/main/config.json

[TensorRT-LLM] TensorRT-LLM version: 0.21.0rc1
[TensorRT-LLM][INFO] Refreshed the MPI local session
Loading weights: 100%|██████████| 541/541 [00:02<00:00, 196.98it/s]
Model init total -- 13.55s
[TensorRT-LLM][INFO] Max KV cache pages per sequence: 257 [window size=8224]
[TensorRT-LLM][INFO] Number of tokens per block: 32.
[TensorRT-LLM][INFO] [MemUsageChange] Allocated 0.44 GiB for max tokens in paged KV cache (8224).
2025-06-16 23:48:52,123 - INFO - flashinfer.jit: Loading JIT ops: norm
2025-06-16 23:48:52,140 - INFO - flashinfer.jit: Finished loading JIT ops: norm
[06/16/2025-23:48:52] [TRT-LLM] [E] Failed to initialize executor on rank 0: rmsnorm(at::Tensor&, at::Tensor&, at::Tensor&, double, bool)::<lambda()> failed to dispatch data type Float
[06/16/2025-23:48:52] [TRT-LLM] [E] Traceback (most recent call last):
  File "/usr/local/lib/python3.12/dist-packages/tensorrt_llm/executor/worker.py", line 713, in worker_main
    worker: GenerationExecutorWorker = worker_cls(
                                       ^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/tensorrt_llm/executor/worker.py", line 139, in __init__
    self.engine = _create_engine()
                  ^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/tensorrt_llm/executor/worker.py", line 137, in _create_engine
    return create_executor(**args)
           ^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/tensorrt_llm/_torch/pyexecutor/py_executor_creator.py", line 333, in create_py_executor
    py_executor = create_py_executor_instance(
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/tensorrt_llm/_torch/pyexecutor/_util.py", line 503, in create_py_executor_instance
    return PyExecutor(resource_manager,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/tensorrt_llm/_torch/pyexecutor/py_executor.py", line 245, in __init__
    self.model_engine.warmup(self.resource_manager)
  File "/usr/local/lib/python3.12/dist-packages/tensorrt_llm/_torch/pyexecutor/model_engine.py", line 493, in wrapper
    return method(self, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/tensorrt_llm/_torch/pyexecutor/model_engine.py", line 747, in warmup
    self.forward(batch,
  File "/usr/local/lib/python3.12/dist-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/tensorrt_llm/_torch/utils.py", line 67, in wrapper
    return func(self, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/tensorrt_llm/_torch/pyexecutor/model_engine.py", line 2058, in forward
    pool = maybe_graph.capture(
           ^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/tensorrt_llm/_torch/pyexecutor/cuda_graph_runner.py", line 110, in capture
    forward_fn(inputs)
  File "/usr/local/lib/python3.12/dist-packages/tensorrt_llm/_torch/pyexecutor/model_engine.py", line 2053, in capture_forward_fn
    return self._forward_step(
           ^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/nvtx/nvtx.py", line 122, in inner
    result = func(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/tensorrt_llm/_torch/pyexecutor/model_engine.py", line 2111, in _forward_step
    logits = self.model_forward(
             ^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/tensorrt_llm/_torch/pyexecutor/model_engine.py", line 2097, in model_forward
    return self.model.forward(**kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/tensorrt_llm/_torch/models/modeling_llama.py", line 920, in forward
    hidden_states = self.model(
                    ^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1762, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/tensorrt_llm/_torch/models/modeling_llama.py", line 860, in forward
    hidden_states, residual = decoder_layer(
                              ^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1762, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/tensorrt_llm/_torch/models/modeling_llama.py", line 572, in forward
    hidden_states = self.input_layernorm(hidden_states)
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1751, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1762, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/tensorrt_llm/_torch/modules/rms_norm.py", line 43, in forward
    hidden_states = flashinfer_rmsnorm(hidden_states, self.weight,
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/_library/custom_ops.py", line 671, in __call__
    return self._opoverload(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/_ops.py", line 756, in __call__
    return self._op(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/_library/custom_ops.py", line 335, in backend_impl
    result = self._backend_fns[device_type](*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/_compile.py", line 51, in inner
    return disable_fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/_dynamo/eval_frame.py", line 838, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/_library/custom_ops.py", line 367, in wrapped_fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/tensorrt_llm/_torch/custom_ops/flashinfer_custom_ops.py", line 47, in flashinfer_rmsnorm
    return rmsnorm(input, weight, eps, enable_pdl=ENABLE_PDL)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/flashinfer/norm.py", line 86, in rmsnorm
    _rmsnorm(out, input, weight, eps, enable_pdl)
  File "/usr/local/lib/python3.12/dist-packages/flashinfer/norm.py", line 98, in _rmsnorm
    get_module_attr("rmsnorm")(out, input, weight, eps, enable_pdl)
  File "/usr/local/lib/python3.12/dist-packages/torch/_ops.py", line 756, in __call__
    return self._op(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: rmsnorm(at::Tensor&, at::Tensor&, at::Tensor&, double, bool)::<lambda()> failed to dispatch data type Float

ERROR:common.base_engine:Failed to start engine: Im

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

• using tag [TensorRT-LLM] TensorRT-LLM version: 0.21.0rc1
• | NVIDIA-SMI 550.90.07 Driver Version: 550.90.07 CUDA Version: 12.8 |

Expected behavior

e.g. https://huggingface.co/baseten/Llama-3.2-3B-Instruct-fp8nokv works fine.

actual behavior

trt-llm crashes with flashinfer error, as the torch_dtype is float32

additional notes

Happy to debug.

Reproduce the failure with TensorRT-LLM 0.21.0rc1 and the linked Orpheus model configurations on the reported environment. Start at _torch/modules/rms_norm.py and _torch/custom_ops/flashinfer_custom_ops.py, then trace the call from _torch/models/modeling_llama.py; done means the model initializes successfully for the reported dtype without the flashinfer dispatch error.