Skip to content

[TRTLLM-5446 Part1] Add torch ref backend - #74

Merged
nvchenghaoz merged 9 commits into
feat/ad-2025-07-07from
chenghao/torch-backend
Jul 10, 2025
Merged

[TRTLLM-5446 Part1] Add torch ref backend#74
nvchenghaoz merged 9 commits into
feat/ad-2025-07-07from
chenghao/torch-backend

Conversation

@nvchenghaoz

Copy link
Copy Markdown

Description

  1. Add a new torch backend. The benefit of new torch backend is facilitate the new feature development effort. And using the torch backend as the reference output in the tests.
  2. Add additional parameters support for the torch backend - soft logitcap / sliding window and attention sinks.
  3. Add necessary tests and update the tests.

Test Coverage

tests/unittest/_torch/auto_deploy/unit/singlegpu/*
tests/unittest/_torch/auto_deploy/unit/singlegpu/transformations/library/test_attention_matcher.py
tests/unittest/_torch/auto_deploy/unit/singlegpu/shim/test_engine.py

Signed-off-by: nvchenghaoz <211069071+nvchenghaoz@users.noreply.github.com>
Signed-off-by: nvchenghaoz <211069071+nvchenghaoz@users.noreply.github.com>
@nvchenghaoz
nvchenghaoz changed the base branch from main to feat/ad-2025-07-07 July 8, 2025 23:10

@lucaslie lucaslie left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Beautiful 🐐

Comment thread tensorrt_llm/_torch/auto_deploy/custom_ops/torch_backend_attention.py Outdated
Comment thread tensorrt_llm/_torch/auto_deploy/custom_ops/torch_backend_attention.py Outdated
Comment thread tensorrt_llm/_torch/auto_deploy/custom_ops/torch_backend_attention.py Outdated
Comment thread tensorrt_llm/_torch/auto_deploy/custom_ops/torch_backend_attention.py Outdated
Signed-off-by: nvchenghaoz <211069071+nvchenghaoz@users.noreply.github.com>
Signed-off-by: nvchenghaoz <211069071+nvchenghaoz@users.noreply.github.com>
Signed-off-by: nvchenghaoz <211069071+nvchenghaoz@users.noreply.github.com>
Signed-off-by: nvchenghaoz <211069071+nvchenghaoz@users.noreply.github.com>

@Fridah-nv Fridah-nv left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thank you! Torch attention backend would be super helpful for debugging.

@lucaslie lucaslie left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Great work 👏

…o/torch-backend

Signed-off-by: nvchenghaoz <211069071+nvchenghaoz@users.noreply.github.com>
…nsorRT-LLM into chenghao/torch-backend

Signed-off-by: nvchenghaoz <211069071+nvchenghaoz@users.noreply.github.com>
Signed-off-by: nvchenghaoz <211069071+nvchenghaoz@users.noreply.github.com>
@nvchenghaoz
nvchenghaoz enabled auto-merge (squash) July 10, 2025 21:07
@nvchenghaoz
nvchenghaoz merged commit 1af3f87 into feat/ad-2025-07-07 Jul 10, 2025
1 check passed
lucaslie pushed a commit that referenced this pull request Jul 18, 2025
* Add the Torch backend and update the test to use the torch backend.

Signed-off-by: nvchenghaoz <211069071+nvchenghaoz@users.noreply.github.com>

* Add the sinks and fix the test failures

Signed-off-by: nvchenghaoz <211069071+nvchenghaoz@users.noreply.github.com>

* address reviewer's comments

Signed-off-by: nvchenghaoz <211069071+nvchenghaoz@users.noreply.github.com>

* use custom op convention

Signed-off-by: nvchenghaoz <211069071+nvchenghaoz@users.noreply.github.com>

* move the ref to the utils_test

Signed-off-by: nvchenghaoz <211069071+nvchenghaoz@users.noreply.github.com>

* Add torch backend tests in ad_build_small_single.py

Signed-off-by: nvchenghaoz <211069071+nvchenghaoz@users.noreply.github.com>

* Address hidden comments...

Signed-off-by: nvchenghaoz <211069071+nvchenghaoz@users.noreply.github.com>

---------

Signed-off-by: nvchenghaoz <211069071+nvchenghaoz@users.noreply.github.com>
Signed-off-by: Lucas Liebenwein <11156568+lucaslie@users.noreply.github.com>
lucaslie pushed a commit that referenced this pull request Jul 21, 2025
* Add the Torch backend and update the test to use the torch backend.

Signed-off-by: nvchenghaoz <211069071+nvchenghaoz@users.noreply.github.com>

* Add the sinks and fix the test failures

Signed-off-by: nvchenghaoz <211069071+nvchenghaoz@users.noreply.github.com>

* address reviewer's comments

Signed-off-by: nvchenghaoz <211069071+nvchenghaoz@users.noreply.github.com>

* use custom op convention

Signed-off-by: nvchenghaoz <211069071+nvchenghaoz@users.noreply.github.com>

* move the ref to the utils_test

Signed-off-by: nvchenghaoz <211069071+nvchenghaoz@users.noreply.github.com>

* Add torch backend tests in ad_build_small_single.py

Signed-off-by: nvchenghaoz <211069071+nvchenghaoz@users.noreply.github.com>

* Address hidden comments...

Signed-off-by: nvchenghaoz <211069071+nvchenghaoz@users.noreply.github.com>

---------

Signed-off-by: nvchenghaoz <211069071+nvchenghaoz@users.noreply.github.com>
Signed-off-by: Lucas Liebenwein <11156568+lucaslie@users.noreply.github.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

4 participants