Use dtype access instead of hardcoding

Author: voltjia

ops/ninetoothed/kernels/scaled_dot_product_attention.py

@@ -30,7 +30,7 @@ def arrange_k_or_v(input):
 
 
 def application(q, k, v, scale, o):
-    q_loaded = (q * scale * 1.44269504089).to(ntl.float16)
+    q_loaded = (q * scale * 1.44269504089).to(q.dtype)
 
     acc = ntl.zeros((q.shape[-2], q.shape[-1]), dtype=ntl.float32)
     l_i = ntl.full((q.shape[-2],), 1, dtype=ntl.float32)
@@ -44,7 +44,7 @@ def application(q, k, v, scale, o):
         l_ij = ntl.sum(p, 1)
 
         alpha = ntl.exp2(m_i - m_ij)
-        acc = acc * alpha[:, None] + ntl.dot(p.to(ntl.float16), v[i])
+        acc = acc * alpha[:, None] + ntl.dot(p.to(v.dtype.dtype), v[i])
         m_i = m_ij
         l_i = l_i * alpha + l_ij