Skip to content

Fix diff_attn_swa plugin int attribute under TRT 10.16 / NumPy 2.x #49

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

Draft
ryanontheinside wants to merge 1 commit into
base: main
Choose a base branch
from
Draft

Fix diff_attn_swa plugin int attribute under TRT 10.16 / NumPy 2.x #49

Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
10 changes: 7 additions & 3 deletions optimized/tensorRT/scripts/diff_attn_nocast_plugin.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,8 @@
"""
import torch
import tensorrt.plugin as trtp
import numpy as np
import numpy.typing as npt
from typing import Tuple

_stream_cache = {}
Expand All @@ -13,15 +15,17 @@

@trtp.register("samel::diff_attn_swa")
def diff_attn_swa_desc(q_bat: trtp.TensorDesc, k_bat: trtp.TensorDesc,
v_bat: trtp.TensorDesc, num_heads: int) -> trtp.TensorDesc:
v_bat: trtp.TensorDesc,
num_heads: npt.NDArray[np.int64]) -> trtp.TensorDesc:
out = q_bat.like()
out.shape_expr[-2] = q_bat.shape_expr[-2] // 2
return out


@trtp.impl("samel::diff_attn_swa")
def diff_attn_swa_impl(q_bat: trtp.Tensor, k_bat: trtp.Tensor, v_bat: trtp.Tensor,
num_heads: int, outputs: Tuple[trtp.Tensor], stream: int):
num_heads: npt.NDArray[np.int64],
outputs: Tuple[trtp.Tensor], stream: int):
global _triton_fn
if stream not in _stream_cache:
_stream_cache[stream] = torch.cuda.ExternalStream(stream)
Expand All @@ -37,5 +41,5 @@ def diff_attn_swa_impl(q_bat: trtp.Tensor, k_bat: trtp.Tensor, v_bat: trtp.Tenso

# NO dtype cast — Triton auto-compiles for the input dtype
o = _triton_fn(q, k, v, window=17)
H = num_heads
H = int(np.asarray(num_heads).reshape(-1)[0])
out_t.copy_(o[:, :, :H, :] - o[:, :, H:, :])