NVIDIA / NVIDIA/TensorRT

Excessive inference time spent on attention_mask operations after ONNX→TensorRT conversion (DriveOS AGX Orin + TRT 8.6.15)

Open
#4,610 2 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Module:Embedded Module:ONNX Module:Performance
Dominant language
C++
Stars
13.4k
Forks
2.4k
Avg merge
5d 3h
Merged PRs (30d)
2

Description


Hi NVIDIA team,

I encountered a severe performance issue when converting a GPT-like model (with attention_mask input) from ONNX to TensorRT.

After successful conversion and deployment on DriveOS (AGX Orin, TensorRT 8.6.15), I used Nsight Systems (nsys) to analyze runtime performance. Surprisingly, more than 90% of total inference time is spent on attention_mask related operations, while actual matrix multiplications and other Transformer layers take very little time.

This behavior seems abnormal — I suspect the mask computation is not being efficiently fused or optimized in TensorRT.


Environment

  • Platform: NVIDIA Drive AGX Orin
  • DriveOS Version: (6.12.1)
  • TensorRT Version: 8.6.15
  • CUDA Version: (11.4)
  • Model type: GPT-like transformer (causal decoder)
  • Input: token embeddings + attention_mask + pask_kv_cache
  • Precision: FP16
  • Batch size: (1)
  • ONNX opset: 16

Steps to Reproduce

  1. Export model from PyTorch → ONNX (with attention_mask input).
  2. Convert ONNX to TensorRT engine:
/usr/src/tensorrt/bin/trtexec --onnx=test.onnx --saveEngine=test.engine --exportLayerInfo=test_layerinfo.log --profilingVerbosity=detailed --exportProfile=test_profile.log --separateProfileRun --duration=5 --streams=1 --useCudaGraph --fp16 --verbose
  1. Run inference and profile using Nsight Systems:
nsys profile --force-overwrite true -o test --trace=cuda,nvtx,cublas,cudnn  --stats=true \
        /usr/src/tensorrt/bin/trtexec  \
        --loadEngine=test.engine  \
       --iterations=10 --idleTime=500 --duration=0 --useSpinWait
  1. Observe runtime results — attention_mask related kernels dominate total time (>90%).

Expected Behavior

attention_mask operations should be lightweight (simple broadcast or add ops) and not dominate runtime.

Due to the attention_mask generating many intermediate (glue) operators, it may be one of the main bottlenecks during TensorRT inference.


Attachments

I’ve attached:

  • test.onnx
  • TensorRT conversion log (test.txt)
  • Nsight Systems trace (test.nsys-rep)

→ Upload as: attention_mask_perf_issue.zip


Questions

  1. Is this a known issue or regression in TRT 8.6.15 on DriveOS?
  2. Are there recommended practices to optimize or fuse attention_mask handling (e.g. using plugin or model rewrite)?

Additional Notes

Additionally, the attached test.onnx is part of the full model. I also attempted removing attention_mask while setting past_key_values_input as a dynamic input. However, after conversion, because of the dynamic input, the latency was still not ideal.

Would really appreciate any feedback or workaround — this issue currently blocks our deployment. Thanks a lot for your time and support!

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with the reported trtexec and nsys commands, then inspect test.onnx, test.txt, and test.nsys-rep if the referenced attachments are available. Compare the attention_mask kernels with the rest of inference and determine whether the result is reproducible; done would require a confirmed cause, supported workaround, or clear explanation of the expected behavior.

Written by the indexing model from the issue text.

Assessment

Tech stack
pytorch
Domain
machine-learning, performance
Issue type
Bug
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.