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

Open
ryanontheinside wants to merge 1 commit into
base: main
Choose a base branch
from
+12 −6
Open

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
18 changes: 12 additions & 6 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 @@ -27,6 +27,8 @@
"""
import os

import numpy as np
import numpy.typing as npt
import torch
import tensorrt.plugin as trtp
from typing import Tuple
Expand Down Expand Up @@ -62,15 +64,17 @@ def _ptx_for(num_heads: int):

@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 @@ -86,13 +90,14 @@ 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:, :])


@trtp.aot_impl("samel::diff_attn_swa")
def diff_attn_swa_aot(q_bat: trtp.TensorDesc, k_bat: trtp.TensorDesc,
v_bat: trtp.TensorDesc, num_heads: int,
v_bat: trtp.TensorDesc,
num_heads: npt.NDArray[np.int64],
outputs: Tuple[trtp.TensorDesc], tactic: int = None
) -> Tuple[str, str, trtp.KernelLaunchParams, trtp.SymIntExprs]:
"""AOT PTX variant — makes the engine graph-capturable. See module docstring.
Expand All @@ -106,14 +111,15 @@ def diff_attn_swa_aot(q_bat: trtp.TensorDesc, k_bat: trtp.TensorDesc,
H2 = q_bat.shape_expr[2]
D = q_bat.shape_expr[3]
H_out = H2 // 2
H = int(np.asarray(num_heads).reshape(-1)[0])

if SWA_AOT_BACKEND == "mma":
# Block-tiled tensor-core kernel. Its scalars are baked in as constexprs at
# compile time, so the only runtime extra is N. Grid is one block per
# (query tile, output head, batch); shared memory holds the Q/K/V tiles.
import swa_mma_aot
from _arch import detect_arch
name, ptx, shared, warps = swa_mma_aot.build(num_heads, detect_arch())
name, ptx, shared, warps = swa_mma_aot.build(H, detect_arch())
extra = trtp.SymIntExprs(1)
extra[0] = N
return (name, ptx,
Expand All @@ -126,7 +132,7 @@ def diff_attn_swa_aot(q_bat: trtp.TensorDesc, k_bat: trtp.TensorDesc,
extra)

# Hand-written scalar kernel: one warp per query, strides passed at runtime.
name, ptx = _ptx_for(num_heads)
name, ptx = _ptx_for(H)
extra = trtp.SymIntExprs(5)
extra[0] = N
extra[1] = H2 * D # stride between positions, input
Expand Down
Loading