From c07698548567fe6f163806f692d282bbaa57aba3 Mon Sep 17 00:00:00 2001
From: RyanOnTheInside <7623207+ryanontheinside@users.noreply.github.com>
Date: Mon, 3 Aug 2026 10:24:39 -0400
Subject: [PATCH] Fix diff_attn_swa integer attributes with the AOT plugin

TensorRT 10.16 supplies ONNX integer attributes as one-element NumPy
arrays. Annotating num_heads as int therefore asks NumPy 2.x to convert
a non-scalar array directly and plugin creation fails before the ONNX can
be parsed.

Keep the ndarray annotation TensorRT expects and explicitly extract the
Python integer inside both the JIT and newer AOT implementations. This
rebases the original fix onto current main after the AOT/MMA plugin work.
---
 .../scripts/diff_attn_nocast_plugin.py         | 18 ++++++++++++------
 1 file changed, 12 insertions(+), 6 deletions(-)

diff --git a/optimized/tensorRT/scripts/diff_attn_nocast_plugin.py b/optimized/tensorRT/scripts/diff_attn_nocast_plugin.py
index 7fe29c5d..1a9cb651 100644
--- a/optimized/tensorRT/scripts/diff_attn_nocast_plugin.py
+++ b/optimized/tensorRT/scripts/diff_attn_nocast_plugin.py
@@ -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
@@ -62,7 +64,8 @@ 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
@@ -70,7 +73,8 @@ def diff_attn_swa_desc(q_bat: trtp.TensorDesc, k_bat: trtp.TensorDesc,
 
 @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)
@@ -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.
@@ -106,6 +111,7 @@ 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
@@ -113,7 +119,7 @@ def diff_attn_swa_aot(q_bat: trtp.TensorDesc, k_bat: trtp.TensorDesc,
         # (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,
@@ -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