Fix operator correctness on Triton and CUDA

Author: whjthu

src/ntops_lab/kernels/creation/affine_grid.py

@@ -12,24 +12,24 @@ def _component_kernel(height, width):
     scale_x = 2.0 / float(width)
 
     def arrangement(theta_row, out_component, scale_y_tensor, scale_x_tensor):
-        theta_arr = theta_row[:, None, None, :].expand((-1, height, width, -1)).flatten(end_dim=3)
+        theta_arr = theta_row[:, None, None, :, :].expand((-1, height, width, -1, -1)).flatten(end_dim=4)
         out_arr = out_component.flatten()
         return theta_arr.tile((1, 3)), out_arr.tile((1,)), scale_y_tensor, scale_x_tensor
 
     def application(theta, out, scale_y_tensor, scale_x_tensor):
         x = ((out.offsets(2) + 0.5) * scale_x_tensor) - 1.0
         y = ((out.offsets(1) + 0.5) * scale_y_tensor) - 1.0
-        t0 = ntl.sum(theta * (theta.offsets(1) == 0), axis=1)
-        t1 = ntl.sum(theta * (theta.offsets(1) == 1), axis=1)
-        t2 = ntl.sum(theta * (theta.offsets(1) == 2), axis=1)
+        t0 = ntl.sum(theta * (theta.offsets(2) == 0), axis=1)
+        t1 = ntl.sum(theta * (theta.offsets(2) == 1), axis=1)
+        t2 = ntl.sum(theta * (theta.offsets(2) == 2), axis=1)
         out = t0 * x + t1 * y + t2
 
     return ninetoothed.make(
         arrangement,
         application,
         (
-            Tensor(2),
             Tensor(3),
+            Tensor(4),
             Tensor(0, constexpr=True, value=scale_y, name="scale_y"),
             Tensor(0, constexpr=True, value=scale_x, name="scale_x"),
         ),
@@ -46,6 +46,5 @@ def run(theta, size, align_corners=False):
     kernel = _component_kernel(height, width)
     scale_y = 2.0 / float(height)
     scale_x = 2.0 / float(width)
-    kernel(theta[:, 0, :], out[..., 0], scale_y, scale_x)
-    kernel(theta[:, 1, :], out[..., 1], scale_y, scale_x)
+    kernel(theta, out, scale_y, scale_x)
     return out

src/ntops_lab/kernels/fused/general/_flash_attention.py

@@ -35,7 +35,7 @@ def application(q, k, v, is_causal, softmax_scale, out):
     m_i = ntl.full((q.shape[-2],), float("-inf"), dtype=ntl.float32)
 
     for i in range(k.shape[0]):
-        qk = ntl.dot(q_loaded, ntl.trans(k[i]))
+        qk = ntl.dot(q_loaded, ntl.trans(k[i])).to(ntl.float32)
         qk = ntl.where(k[i].offsets(-2) < k.source.shape[-2], qk, float("-inf"))
 
         if is_causal:

src/ntops_lab/kernels/fused/general/add_layer_norm_sigmoid.py

@@ -16,7 +16,7 @@ def application(x, residual, gamma, beta, out, hidden):
     centered = y - mean[:, None]
     var = ntl.sum(centered * centered, axis=1) / hidden
     value = centered * ntl.rsqrt(var[:, None] + 1.0e-5) * gamma + beta
-    out = ntl.sigmoid(value)
+    out = (1.0 / (1.0 + ntl.exp(-(value))))
 
 @functools.cache
 def _kernel(hidden):

src/ntops_lab/kernels/fused/general/add_layer_norm_silu.py

@@ -16,7 +16,7 @@ def application(x, residual, gamma, beta, out, hidden):
     centered = y - mean[:, None]
     var = ntl.sum(centered * centered, axis=1) / hidden
     value = centered * ntl.rsqrt(var[:, None] + 1.0e-5) * gamma + beta
-    out = value * ntl.sigmoid(value)
+    out = value * (1.0 / (1.0 + ntl.exp(-(value))))
 
 @functools.cache
 def _kernel(hidden):

src/ntops_lab/kernels/fused/general/add_rms_norm_sigmoid.py

@@ -13,7 +13,7 @@ def application(x, residual, weight, out, hidden):
     y = x + residual
     rrms = ntl.rsqrt(ntl.sum(y * y, axis=1) / hidden + 1.0e-5)
     value = y * rrms[:, None] * weight
-    out = ntl.sigmoid(value)
+    out = (1.0 / (1.0 + ntl.exp(-(value))))
 
 @functools.cache
 def _kernel(hidden):

src/ntops_lab/kernels/fused/general/add_rms_norm_silu.py

@@ -13,7 +13,7 @@ def application(x, residual, weight, out, hidden):
     y = x + residual
     rrms = ntl.rsqrt(ntl.sum(y * y, axis=1) / hidden + 1.0e-5)
     value = y * rrms[:, None] * weight
-    out = value * ntl.sigmoid(value)
+    out = value * (1.0 / (1.0 + ntl.exp(-(value))))
 
 @functools.cache
 def _kernel(hidden):

src/ntops_lab/kernels/fused/general/fused_addmm_bias_abs.py

@@ -3,9 +3,9 @@
 import ninetoothed.language as ntl
 from ninetoothed import Tensor, block_size
 
-BM = block_size()
-BN = block_size()
-BK = block_size()
+BM = block_size(upper_bound=64)
+BN = block_size(upper_bound=64)
+BK = block_size(upper_bound=64)
 
 def arrangement(c, a, b, bias, out):
     out_arr = out.tile((BM, BN))
@@ -14,7 +14,7 @@ def arrangement(c, a, b, bias, out):
     a_arr.dtype = a_arr.dtype.squeeze(0)
     b_arr = b.tile((BK, BN)).tile((-1, 1)).expand((out_arr.shape[0], -1))
     b_arr.dtype = b_arr.dtype.squeeze(1)
-    bias_arr = bias.tile((BN,)).unsqueeze(0).expand((out_arr.shape[0], -1))
+    bias_arr = bias[None, :].expand((out.shape[0], -1)).tile((BM, BN))
     return c_arr, a_arr, b_arr, bias_arr, out_arr
 
 def application(c, a, b, bias, out):

src/ntops_lab/kernels/fused/general/fused_addmm_bias_clamp.py

@@ -3,9 +3,9 @@
 import ninetoothed.language as ntl
 from ninetoothed import Tensor, block_size
 
-BM = block_size()
-BN = block_size()
-BK = block_size()
+BM = block_size(upper_bound=64)
+BN = block_size(upper_bound=64)
+BK = block_size(upper_bound=64)
 
 def arrangement(c, a, b, bias, out):
     out_arr = out.tile((BM, BN))
@@ -14,7 +14,7 @@ def arrangement(c, a, b, bias, out):
     a_arr.dtype = a_arr.dtype.squeeze(0)
     b_arr = b.tile((BK, BN)).tile((-1, 1)).expand((out_arr.shape[0], -1))
     b_arr.dtype = b_arr.dtype.squeeze(1)
-    bias_arr = bias.tile((BN,)).unsqueeze(0).expand((out_arr.shape[0], -1))
+    bias_arr = bias[None, :].expand((out.shape[0], -1)).tile((BM, BN))
     return c_arr, a_arr, b_arr, bias_arr, out_arr
 
 def application(c, a, b, bias, out):

src/ntops_lab/kernels/fused/general/fused_addmm_bias_cube.py

@@ -3,9 +3,9 @@
 import ninetoothed.language as ntl
 from ninetoothed import Tensor, block_size
 
-BM = block_size()
-BN = block_size()
-BK = block_size()
+BM = block_size(upper_bound=64)
+BN = block_size(upper_bound=64)
+BK = block_size(upper_bound=64)
 
 def arrangement(c, a, b, bias, out):
     out_arr = out.tile((BM, BN))
@@ -14,7 +14,7 @@ def arrangement(c, a, b, bias, out):
     a_arr.dtype = a_arr.dtype.squeeze(0)
     b_arr = b.tile((BK, BN)).tile((-1, 1)).expand((out_arr.shape[0], -1))
     b_arr.dtype = b_arr.dtype.squeeze(1)
-    bias_arr = bias.tile((BN,)).unsqueeze(0).expand((out_arr.shape[0], -1))
+    bias_arr = bias[None, :].expand((out.shape[0], -1)).tile((BM, BN))
     return c_arr, a_arr, b_arr, bias_arr, out_arr
 
 def application(c, a, b, bias, out):

src/ntops_lab/kernels/fused/general/fused_addmm_bias_elu.py

@@ -3,9 +3,9 @@
 import ninetoothed.language as ntl
 from ninetoothed import Tensor, block_size
 
-BM = block_size()
-BN = block_size()
-BK = block_size()
+BM = block_size(upper_bound=64)
+BN = block_size(upper_bound=64)
+BK = block_size(upper_bound=64)
 
 def arrangement(c, a, b, bias, out):
     out_arr = out.tile((BM, BN))
@@ -14,7 +14,7 @@ def arrangement(c, a, b, bias, out):
     a_arr.dtype = a_arr.dtype.squeeze(0)
     b_arr = b.tile((BK, BN)).tile((-1, 1)).expand((out_arr.shape[0], -1))
     b_arr.dtype = b_arr.dtype.squeeze(1)
-    bias_arr = bias.tile((BN,)).unsqueeze(0).expand((out_arr.shape[0], -1))
+    bias_arr = bias[None, :].expand((out.shape[0], -1)).tile((BM, BN))
     return c_arr, a_arr, b_arr, bias_arr, out_arr
 
 def application(c, a, b, bias, out):