Fix the boundary issues

Author: voltjia

ops/ninetoothed/kernels/scaled_dot_product_attention.py

@@ -38,6 +38,7 @@ def application(q, k, v, scale, o):
 
     for i in range(k.shape[0]):
         qk = ntl.dot(q_loaded, ntl.trans(k[i]))
+        qk = ntl.where(k[i].offsets(-2) < k.source.shape[-2], qk, float("-inf"))
 
         m_ij = ntl.maximum(m_i, ntl.max(qk, 1))
         p = ntl.exp2(qk - m_ij[:, None])
@@ -49,7 +50,7 @@ def application(q, k, v, scale, o):
         l_i = l_i * alpha + l_ij
 
     acc /= l_i[:, None]
-    o = acc  # noqa: F841
+    o = acc.to(o.dtype)  # noqa: F841
 
 
 shape_options = (None, None, None, {"constexpr": True, "upper_bound": 128})

ops/triton/kernels/scaled_dot_product_attention.py

@@ -96,23 +96,24 @@ def kernel(
         order=(1, 0),
     )
 
-    q = (tl.load(q_block_ptr) * scale * 1.44269504089).to(q_block_ptr.type.element_ty)
+    q = tl.load(q_block_ptr, boundary_check=(0, 1))
+    q = (q * scale * 1.44269504089).to(q_block_ptr.type.element_ty)
 
     acc = tl.zeros((BLOCK_SIZE_M, EMB_DIM), dtype=tl.float32)
     l_i = tl.full((BLOCK_SIZE_M,), 1, dtype=tl.float32)
     m_i = tl.full((BLOCK_SIZE_M,), float("-inf"), dtype=tl.float32)
 
-    for _ in range(0, tl.cdiv(seq_len, BLOCK_SIZE_N)):
-        k = tl.load(k_block_ptr)
+    for i in range(0, tl.cdiv(seq_len, BLOCK_SIZE_N)):
+        k = tl.load(k_block_ptr, boundary_check=(0, 1))
 
-        qk = tl.dot(q, k)
+        mask = i * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) < seq_len
+        qk = tl.where(mask, tl.dot(q, k), float("-inf"))
 
         m_ij = tl.maximum(m_i, tl.max(qk, 1))
-        qk -= m_ij[:, None]
-        p = tl.exp2(qk)
+        p = tl.exp2(qk - m_ij[:, None])
         l_ij = tl.sum(p, 1)
 
-        v = tl.load(v_block_ptr)
+        v = tl.load(v_block_ptr, boundary_check=(0, 1))
         alpha = tl.exp2(m_i - m_ij)
         acc = acc * alpha[:, None] + tl.dot(p.to(v_block_ptr.type.element_ty), v)
         m_i = m_ij
@@ -123,4 +124,4 @@ def kernel(
 
     acc /= l_i[:, None]
 
-    tl.store(o_block_ptr, acc.to(o_ptr.type.element_ty))
+    tl.store(o_block_ptr, acc.to(o_ptr.type.element_ty), boundary_check=(0, 1))

scaled_dot_product_attention.py

@@ -109,7 +109,7 @@ def _rope(x, sin_table, cos_table):
         print("✅ NineToothed and PyTorch match.")
     else:
         print("❌ NineToothed and PyTorch differ.")
-    if torch.allclose(ninetoothed_output, triton_output, atol=0.001, rtol=0.001):
+    if torch.allclose(ninetoothed_output, triton_output, atol=0, rtol=0):
         print("✅ NineToothed and Triton match.")
     else:
         print("❌ NineToothed and Triton differ.")
@@ -143,7 +143,7 @@ def benchmark(seq_len, provider):
         triton_output = ops.triton.torch.scaled_dot_product_attention(q, k, v)
 
         assert torch.allclose(ninetoothed_output, torch_output, atol=0.025, rtol=0.025)
-        assert torch.allclose(ninetoothed_output, triton_output, atol=0.001, rtol=0.001)
+        assert torch.allclose(ninetoothed_output, triton_output, atol=0, rtol=0)
 
         if provider == "ninetoothed":
             ms = triton.testing.do_bench(