Skip to content

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

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
nvchenghaoz merged 9 commits into from
Jul 10, 2025
Merged

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

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
9 commits
Select commit Hold shift + click to select a range
cae02e2
Add the Torch backend and update the test to use the torch backend.
nvchenghaoz Jul 1, 2025
c694029
Add the sinks and fix the test failures
nvchenghaoz Jul 8, 2025
e7b965d
address reviewer's comments
nvchenghaoz Jul 9, 2025
5ac2b08
use custom op convention
nvchenghaoz Jul 9, 2025
bffda56
move the ref to the utils_test
nvchenghaoz Jul 9, 2025
bc8a6ab
Add torch backend tests in ad_build_small_single.py
nvchenghaoz Jul 9, 2025
82e3443
Merge remote-tracking branch 'origin/feat/ad-2025-07-07' into chengha…
nvchenghaoz Jul 10, 2025
1d5d554
Merge branch 'chenghao/torch-backend' of github.com:nv-auto-deploy/Te…
nvchenghaoz Jul 10, 2025
9740d0b
Address hidden comments...
nvchenghaoz Jul 10, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions tensorrt_llm/_torch/auto_deploy/custom_ops/__init__.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,7 @@
from .mla import *
from .quant import *
from .torch_attention import *
from .torch_backend_attention import *
from .torch_moe import *
from .torch_rope import *
from .triton_attention import *
Expand Down
2 changes: 2 additions & 0 deletions tensorrt_llm/_torch/auto_deploy/custom_ops/torch_attention.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,8 @@
import torch.nn as nn
import torch.nn.functional as F

# TODO (chenghao): Remove related kernels once we have a backend-specific implementation for attention.


@torch.library.custom_op("auto_deploy::torch_attention_repeat_kv", mutates_args=())
def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:
Expand Down
Loading